torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

Codeforces #271, B : Worms

問題概要

 ユニークな通し番号を付けた虫たちを n 個の塊に分ける。結果、i 番目 ( 1-indexed ) の塊には a_i 匹の虫がいて、その通し番号は( 0 番目の山にいる虫の数を便宜的に a_0 とおけば)( \sum_{j = 0}^{i - 1}a_j,  ( \sum_{j = 0}^{i - 1}a_j ) + a_i ] である。
 以下の形式のクエリを m 件処理せよ。

  • 整数 q が与えられる。通し番号 q をもつ虫がいる塊の番号を出力せよ。

 |a| \leq 10^5
 m \leq 10^5

解法

 もしクエリの数が少なければ、塊を先頭から走査して答えを求めることができますが、それでは TLE します。より高速にクエリを処理する必要があります。
 a の累積和をとったものを b とします。すなわち、b_i = \sum_{ j = 1 }^{i}a_j です。このとき、クエリ q に対する答えは、b_r \leq q である最大の r となります。b_i の値は i に関して単調性をもつので、二分探索を適用することができます。こうすることで、前処理に O( n ) 時間、全クエリを O( m \log n ) 時間で処理することができます。

コード

using VI = vector<int>; using VVI = vector<VI>;
template < typename T = int > using VT = vector<T>;

template < typename T > inline void input_n( VT< T > &out ) { copy_n( istream_iterator< T >( cin ), out.size(), out.begin() ); };

#define FOR( e, c ) for ( auto &e : c )
#define ALL( c ) (c).begin(), (c).end()

int main()
{
	cin.tie( 0 );
	ios::sync_with_stdio( false );

	int n;
	cin >> n;

	VI as( n );
	input_n( as );

	int m;
	cin >> m;

	VI qs( m );
	input_n( qs );

	VI csum;
	partial_sum( ALL( as ), back_inserter( csum ) );

	FOR( q, qs )
	{
		cout << lower_bound( ALL( csum ), q ) - csum.begin() + 1 << endl;
	}

	return 0;
}