AtCoder ARC #014 : D - grepマスター

さてD問題。今回はCがハマる問題な上、Dが楽なので正答者が多めだったね。
自分もCよりDの方がミスなく早く回答できている。
http://arc014.contest.atcoder.jp/tasks/arc014_4

問題

A行の文章に"grep -B x -A y 'needle'"というコマンドを実行することを考える。
こうすると文章中'needle'という単語を含む行と、その前x行とその後y行が出力される。
ただし、複数のneedleを含む行同士が近くにある場合でも、その間の行は重複して出力はされない。

needleという単語を含む要素数Nの行番号の配列Lと、(x,y)の組からなるクエリがM個与えられる。
各クエリについて、出力される行数を答えよ。

解法

A<=10^9なので、どの行が出力されるかを律儀に計算する時間はない。
もっと早くクエリに答える手段を考える。

各クエリの出力行数は以下の和である。

• x,yの値を問わず、needleを含むN行は必ず出力される。
• １つ目のneedle行L[1]の前x行は出力される。ただしL[1]<=xならL[1]-1行だけ
• 最後ののneedle行L[N]の後y行は出力される。ただしA-L[N]
• 各L[i]とL[i+1]の間において、(x+y)行が出力される。ただし両者の間の行数(L[i+1]-L[i]-1)が(x+y)未満ならその行数が出力される。

先頭3つは各クエリに対しO(1)で処理できる。
問題は最後で、このまま処理するとO(N)かかる。
このままだとクエリ数と合わせると全体でO(NM)かかるのでTLEしてしまう。

そこで最後の処理を高速化する。
最後の処理は行数だけが問題で行の位置は関係ないので、間の行数(L[i+1]-L[i]-1)を事前に計算して配列にしまい、ソートしておく。
各クエリに対し、間の行数が(x+y)以上の要素数分については(x+y)×( (x+y)行以上の要素数)を加算し、(x+y)未満の要素については、それらの和を加算すればよい。
(x+y)未満の要素の和は、先に累積和を取り出しておく。

こうすると、クエリに対し(x+y)以上の要素数を求めるのがO(logN)でできるので、全クエリをO(M logN)で処理できる。

ll A,N,M;
vector<ll> L;
vector<ll> K,K2;

void solve() {
	int f,r,i,j,k,l;
	ll x,y;
	
	cin>>A>>N>>M;
	L.resize(N);
	FOR(i,N) cin>>L[i];
	FOR(i,N-1) K.push_back(L[i+1]-L[i]-1);
	sort(K.begin(),K.end());
	
	K2.push_back(0);
	FOR(i,K.size()) K2.push_back(K2[i]+K[i]);
	
	FOR(i,M) {
		cin >> x >> y;
		ll ret=N;
		
		ret += min(L[0]-1,x);
		ret += min(A-L[N-1],y);
		x=x+y;
		
		vector<ll>::iterator vit=upper_bound(K.begin(),K.end(),x);
		y=vit-K.begin();
		ret += K2[y] + (N-1-y)*(ll)x;
		cout << ret << endl;
	}
	return;
}

まとめ

いつもlower_boundとupper_boundが迷うんだよなぁ。
こちらはノーミス１発で解けて良かった。