Codeforces #202 Div1. C. Subset Sums

こちらはCodeforcesっぽい問題。
http://codeforces.com/contest/348/problem/C

問題

N要素からなる整数列Aと、1～Nの値からなる集合M個からなる配列Sが与えられる。
これらのデータに対し、以下のクエリをQ回実行せよ。

1. iが与えられるので、S[i]に含まれる整数値S[i][j]に対し、A[S[i][j]]の総和を答える
2. iとxが与えられるので、S[i]に含まれる整数値S[i][j]に対し、A[S[i][j]]にxを加算する。

なお、Sの要素に含まれる値の合計は10^5以下である。

解法

2つ目のクエリに対して、Aの要素を逐一更新する手法と、Sの総和を逐一更新する手法がある。
しかしそれぞれO(NQ)またはO(MQ)かかってしまい、NかMが多いとTLEする。

ここで、Sの要素に含まれる値の総和は小さいことを利用する。
まず、S[i]の中身が多いものと少ないものに分ける。
Sの要素の総和は10^5しかないので、中身が多いS[i]はそこまで多くない。

2つ目のクエリに対し、S[i]の要素が多い場合はAをいちいち更新せず、他のS[j]とのintersectionの数だけxを足しこむことで高速化する。

まず、中身が多いS[i]と、全てのS[j]に対しintersectionの数を求めておく。
また、初期値として各S[i]についてA[S[i][j]]の総和を1度計算しておく。

• クエリ1に対し：
  • S[i]の中身が多いなら、計算済みの総和を出す。
  • S[i]の中身が少ないなら、A[S[i][j]]を計算し、かつ中身の多いS[j]とのintersectionの要素数と、S[j]に加算済みの値を掛けたものを足す。
• クエリ2に対し：
  • S[i]の中身が多いなら、Aをいちいち更新せず、S[i]の総和にxを加えておく。
  • S[i]の中身が少ないなら、Aを更新する。
  • いずれの場合も、中身の多いS[j]に対してはS[i]とS[j]とのintersectionの要素数とxの積を足す。

int N,M,Q;
ll A[100011],ST[100011];
ll AC[100011];
vector<int> S[100001];
vector<int> li,he;
map<ll,int> intersection;

void solve() {
	int f,i,j,k,x,y,z,tx,ty;
	
	cin>>N>>M>>Q;
	FOR(i,N) cin>>A[i];
	FOR(i,M) {
		cin>>x;
		FOR(j,x) {
			cin>>y;
			y--;
			ST[i]+=A[y];
			S[i].push_back(y);
		}
		sort(S[i].begin(),S[i].end());
	}
	FOR(i,M) {
		if(S[i].size()<2000) li.push_back(i);
		else he.push_back(i);
	}
	
	FOR(x,N) FOR(y,he.size()) {
		vector<int> vv;
		set_intersection(S[x].begin(),S[x].end(),S[he[y]].begin(),S[he[y]].end(),back_inserter(vv));
		intersection[x*100000LL+he[y]]=vv.size();
	}
	
	FOR(i,Q) {
		string s;
		cin>>s>>k;
		k--;
		if(s[0]=='?') {
			ll r = 0;
			if(S[k].size()<2000) {
				FOR(j,S[k].size()) r+=A[S[k][j]];
				FOR(y,he.size()) r+= intersection[k*100000LL+he[y]]*AC[he[y]];
			}
			else r=ST[k];
			
			cout << r << endl;
		}
		else {
			cin>>x;
			if(S[k].size()<2000) FOR(j,S[k].size()) A[S[k][j]]+=x;
			else AC[k]+=x;
			FOR(y,he.size()) ST[he[y]] += intersection[k*100000LL+he[y]]*(ll)x;
		}
	}
}

まとめ

これ、S[i]の大きい・少ないの基準を小さくしすぎるとTLEした。