kmjp's blog

競技プログラミング参加記です

Thanks Kosen 2020 : G - 答辞

誕生日で数列はよく見る設定だけど、答辞は珍しい。
https://www.hackerrank.com/contests/thankskosen2020/challenges/challenge-2436

問題

N要素の整数列A,Bが与えられる。
区間[L,R]のうち、B[L]≦sum(A[L...R])≦B[R]となるものを求めよ。

Lをずらしつつ、sum(A[L...R])≦B[R]となるRを数え上げていこう。
まず各Rについて、Lをだんだん大きくしたときにはじめてsum(A[L...R])≦B[R]となるタイミングを考えよう。
これはsum(A[0...R])-B[R]をsetに放り込んでおいて、sum(A[0...(L-1)])がsum(A[0...R])-B[R]を超えるタイミングを見ればよい。
この時点で、Rが選択肢となりうるのでBITでR番目の要素を1インクリメントしておく。
こうすればLに対するRを数え上げられる。

一方B[L]≦sum(A[L...R])の条件も満たさなければならないが、変形するとB[L]≦sum(A[0...R])-sum(A[0...(L-1)])なのでB[L]+sum(A[0...(L-1)])≦sum(A[0...R])を満たすRが対象となる。
これはAの累積和を取っておけば条件を満たす最小のRが二分探索で求められる。

これでLに対しB[L]≦sum(A[L...R])を満たすRの範囲を求め、sum(A[L...R])≦B[R]を満たす要素数をBITで数え上げられるようになる。

int N;
ll A[101010],B[101010],S[101010];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<int,20> bt;

void solve() {
	int i,j,k,l,r,x,y;;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i];
		S[i]=A[i];
		if(i) S[i]+=S[i-1];
	}
	FOR(i,N) {
		cin>>B[i];
	}
	set<pair<ll,int>> C;
	FOR(i,N) {
		if(S[i]<=B[i]) {
			bt.add(i,1);
		}
		else {
			C.insert({S[i]-B[i],i});
		}
	}
	
	ll ret=0;
	ll s=0;
	FOR(i,N) {
		x=lower_bound(S,S+N,B[i]+(i?S[i-1]:0))-S;
		x=max(x,i);
		ret+=bt(N)-bt(x-1);
		s+=A[i];
		while(C.size() && C.begin()->first<=s) {
			bt.add(C.begin()->second,1);
			C.erase(C.begin());
		}
	}
	cout<<ret<<endl;
	
	
}

まとめ

数値を200000個も読み上げる答辞はやる方も聞く方も勘弁。