Codeforces #890 : Div2 E2. PermuTree (hard version)

途中C問題で躓いてひどいことに。
https://codeforces.com/contest/1856/problem/E2

問題

N頂点の根付き木が与えられる。
各点iに1～Nの値を1つずつ付与した値をA[i]とする。
最適なAの割り当てをした時、頂点対(u,v)に対しA[u]＜A[LCA(u,v)]＜A[v]となるような(u,v)の数は何通りか。

解法

各点iに対し、そのSubtree群をサイズの和が極力均等になるように2分割する。
片方にA[i]未満、片方にA[i]より大きな値を割り当てればよい。

あとはいかに極力均等になるようにするかだが、基本的にBitsetでDPする。
頂点毎に毎回N bitのbitsetを使う余裕はないので、Nのサイズに応じたサイズのBitsetを使うよう、テンプレートをうまく使おう。

int N;
vector<int> E[1010101];
int P[1050501];
int C[1010101];

unordered_map<int,int> M;

template <int len = 64>
ll hoge(vector<int>& V) {
	if(V.size()<=1) return 0;
	ll sum=0;
	int i;
	
	int mv=0;
	
	M.clear();
	FORR(v,V) {
		mv=max(mv,v);
		sum+=v;
		M[v]++;
	}
	if(mv*2>=sum) {
		return 1LL*mv*(sum-mv);
	}
	
	if(sum+1>len) return hoge<min(1<<21,len*2)>(V);
	
	bitset<len> B;
	B[0]=1;
	ll ma=0;
	FORR2(a,b,M) {
		int x=1;
		b++;
		while(x*2<=b) {
			B|=B<<(x*a);
			x*=2;
		}
		B|=B<<((b-x)*a);
	}
	
	FOR(i,sum+1) if(B[i]) ma=max(ma,1LL*i*(sum-i));
	
	return ma;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	for(i=1;i<N;i++) {
		cin>>P[i];
		//P[i]=1+rand()%i;
		P[i]--;
		E[P[i]].push_back(i);
	}
	
	ll ret=0;
	for(i=N-1;i>=0;i--) {
		C[i]=1;
		vector<int> V;
		FORR(e,E[i]) {
			C[i]+=C[e];
			V.push_back(C[e]);
		}
		ret+=hoge(V);
	}
	cout<<ret<<endl;
}

まとめ

Bitsetのサイズを動的に切り替えるこのトリック、初めて見た。