kmjp's blog

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

Codeforces #489 Div2 E. Nastya and King-Shamans

ゴリ押しですね。
http://codeforces.com/contest/992/problem/E

問題

整数列Aが与えられる。
その後、Q個のクエリが与えられる。
各クエリでは、Aの要素を1つ更新する。

クエリを更新するごとに、Aのうちprefixの和と自身の値が等しい要素が存在すればそれを1つ答えよ。
言い換えるとsum(A[0...(i-1)])=A[i]となるiを1つ答えよ。

解法

平方分割で解く。
バケット毎に、(A[i]-(バケット内で自分の手前にある値の和))をキーとし、要素番号をvalueとしたmapを作る。
値の更新クエリの際には、1バケット分だけmapを作り直そう。

条件を満たす要素の検索は、バケット毎にprefixの総和を求めつつ、各バケットの要素を検索すれば解が出る。
なお、mapを使うとTLEするのでunordered_mapを使おう。

int N,Q;
int A[202020];
const int DI=450;
unordered_map<ll,int> S[DI];
ll tsum[450];

void update(int id) {
	int j;
	ll sum=0;
	S[id].clear();
	FOR(j,DI) {
		if(id*DI+j>=N) break;
		if(A[id*DI+j]>=sum) S[id][A[id*DI+j]-sum]=id*DI+j;
		sum+=A[id*DI+j];
	}
	tsum[id]=sum;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N) cin>>A[i];
	FOR(i,DI) update(i);
	
	while(Q--) {
		cin>>x>>y;
		x--;
		A[x]=y;
		update(x/DI);
		
		ll sum=0;
		FOR(i,DI) {
			if(S[i].empty()) continue;
			if(S[i].count(sum)) {
				cout<<S[i][sum]+1<<endl;
				break;
			}
			sum+=tsum[i];
		}
		
		if(i==DI) cout<<-1<<endl;
		
	}
	
	
}

まとめ

unordered_mapでゴリ押した感。