 F. Unique Occurrences 线段树分治)
F. Unique Occurrences
time limit per test: 6 seconds
memory limit per test: 1024 megabytes
input: standard input
output: standard output
You are given a tree, consisting ofnnnvertices. Each edge has an integer value written on it.
Letf(v,u)f(v, u)f(v,u)be the number of values that appearexactly onceon the edges of a simple path between verticesvvvanduuu.
Calculate the sum off(v,u)f(v, u)f(v,u)over all pairs of verticesvvvanduuusuch that1≤v<u≤n1 \le v < u \le n1≤v<u≤n.
DeepL 翻译:
给你一棵由nnn个顶点组成的树。每条边上都写有一个整数值。
假设f(v,u)f(v,u)f(v,u)是顶点vvv和uuu之间简单路径的边上只出现一次的值的个数。
计算所有顶点vvv和uuu之间的f(v,u)f(v,u)f(v,u)和1≤v<u≤n1 ≤ v < u ≤ n1≤v<u≤n。
把每种颜色看为时间轴上的事件，进行分治处理。
简单来讲，我们用线段树分治，对于每种颜色全部删除后的树上的 dsu 进行统计，很显然，此时某一条边的贡献就是 左右两个连通块的大小的乘积，然后再把这种颜色的边加回来，再进行下一条边的枚举。
总的复杂度只有 n log n
vector<tuple<int,int,int>>e;vector<vector<PII>>col;intans;structop{intt,a,b;};classDSU{public:vector<int>f,siz;stack<op>save;DSU(){}DSU(intn){init(n);}voidinit(intn){f.resize(n+1);iota(range(f),0);siz.assign(n+1,1);}intfind(intx){returnf[x]==x?x:find(f[x]);}boolsame(intx,inty){returnfind(x)==find(y);}boolmerge(intx,inty){x=find(x);y=find(y);if(x==y){return0;}if(siz[x]>siz[y]){swap(x,y);}save.push({1,x,y});siz[y]+=siz[x];f[x]=y;returntrue;}voidrollback(){auto[t,x,y]=save.top();save.pop();if(t==1){siz[y]-=siz[x];f[x]=x;}else{}}size_tgetSnapshot()const{returnsave.size();}voidroll(inttar){while(save.size()>tar){rollback();}}intsize(intx){returnsiz[find(x)];}};template<typenameT>classSegmentTreeDC{public:inttot;// 时间点数（叶子数量）vector<vector<T>>tr;// 每个节点的操作列表SegmentTreeDC():tot(0LL){}SegmentTreeDC(int_T){init(_T);}voidinit(int_T){tot=max(1LL,_T);tr.clear();tr.resize(4LL*(tot+5));}// 插入一个操作 op 到闭区间 [l, r]voidinsert(intp,intL,intR,intl,intr,constT&op){if(l>r||l>R||r<L)return;if(l<=L&&R<=r){tr[p].push_back(op);return;}intmid=L+R>>1;insert(p<<1,L,mid,l,r,op);insert(p<<1|1,mid+1,R,l,r,op);}voidinsert(intl,intr,constT&op){if(l>r)return;insert(1,1,tot,l,r,op);}voiddfs(DSU&dsu){function<void(int,int,int)>dfs=[&](intp,intl,intr){size_t snap=dsu.getSnapshot();for(auto&e:tr[p]){dsu.merge(e.first,e.second);}if(l==r){for(auto[u,v]:col[l]){ans+=dsu.size(u)*dsu.size(v);}}else{intmid=(l+r)>>1;dfs(p<<1,l,mid);dfs(p<<1|1,mid+1,r);}dsu.roll(snap);};dfs(1,1,tot);}};voidsolve(){intn;cin>>n;e.resize(n+1);col.resize(n+2);for(inti=1;i<n;i++){intu,v,w;cin>>u>>v>>w;e[i]={u,v,w};col[w].push_back({u,v});}DSUdsu(n);SegmentTreeDC<PII>seg(n);for(inti=1;i<n;i++){auto&[u,v,w]=e[i];seg.insert(1,w-1,{u,v});seg.insert(w+1,n,{u,v});}seg.dfs(dsu);cout<<ans<<endl;}signedmain(){cin.tie(nullptr)->ios::sync_with_stdio(false);intT=1;// cin >> T;for(;T_<=T;T_++)solve();return0;}
