Codeforces ECR #129 : F. Unique Occurrences

6問回にしては軽め？
https://codeforces.com/contest/1681/problem/F

問題

木を成す無向グラフが与えられる。
各辺には整数値が設定されている。

f(u,v) := パスu-v間の辺に設定された整数値のうち、1回しか登場しないものの数
としたとき、u＜vの範囲でu,vを総当たりしたときのf(u,v)の総和を求めよ。

解法

辺の値x毎に考える。
xの辺以外を縮約したグラフで、xの辺で隣接する連結成分のサイズの積を取ればよい。
BITを使い、連結成分毎に隣接する連結成分のサイズを数え上げよう。

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,rem;

int N;
vector<int> E[505050];
int L[505050],R[505050];
int PE[505050];
int U[505050],V[505050],X[505050];
vector<int> Xs[505050];

int id;
void dfs(int cur,int pre) {
	L[cur]=id++;
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur);
	R[cur]=id;
}

vector<pair<int,int>> cand;
vector<pair<int,int>> hist,hist2;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>U[i]>>V[i]>>X[i];
		U[i]--,V[i]--,X[i]--;
		E[U[i]].push_back(V[i]);
		E[V[i]].push_back(U[i]);
		Xs[X[i]].push_back(i);
	}
	dfs(0,0);
	FOR(i,N) bt.add(i,1);
	
	ll ret=0;
	FOR(i,N) if(Xs[i].size()) {
		cand.clear();
		hist.clear();
		hist2.clear();
		FORR(e,Xs[i]) {
			if(L[U[e]]<L[V[e]]) cand.push_back({L[V[e]],V[e]});
			else cand.push_back({L[U[e]],U[e]});
		}
		sort(ALL(cand));
		reverse(ALL(cand));
		cand.push_back({0,0});
		FORR2(l,v,cand) {
			ll a=bt(R[v]-1)-bt(L[v]-1);
			ll b=rem(R[v]-1)-rem(L[v]-1);
			ret+=a*b;
			if(v) {
				bt.add(L[v],-a);
				rem.add(L[v],a-b);
				hist.push_back({L[v],a});
				hist2.push_back({L[v],-a+b});
			}
		}
		FORR2(a,b,hist) bt.add(a,b);
		FORR2(a,b,hist2) rem.add(a,b);
	}
	cout<<ret<<endl;
}

まとめ

もっとすんなり書けるといいんだけどな。