C133 线段树分治+并查集 CF1681F Unique Occurrences

视频链接：C133 线段树分治 CF1681F Unique Occurrences_哔哩哔哩_bilibili

 

 

Unique Occurrences - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

Problem - 1681F- Codeforces

// 线段树分治 O(nlognlogn)
#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;

#define int long long
#define ls (u<<1)
#define rs (u<<1|1)
#define mid (l+r>>1)
const int N=500005;

int n,m,top,ans;
int p[N],high[N],siz[N];
vector<pair<int,int> >e[N]; //存边
vector<pair<int,int> >tr[N<<2];
struct node{
  int x,y,hy;
}st[N]; //用栈记录边的端点的合并信息

void insert(int u,int l,int r,int L,int R,pair<int,int>e){
  if(L<=l&&r<=R){
    tr[u].push_back(e); //把颜色边挂在出现时间段的节点上
    return;
  }
  if(L<=mid) insert(ls,l,mid,L,R,e);
  if(R>mid) insert(rs,mid+1,r,L,R,e);
}
int find(int x){ //查找根
  return p[x]==x?x:find(p[x]);
}
void merge(int x,int y){ //合并集合
  x=find(x),y=find(y);
  if(x==y) return;
  if(high[x]>high[y])swap(x,y);
  st[++top]={x,y,high[y]};
  p[x]=y; //x连向y
  high[y]+=(high[y]==high[x]);
  siz[y]+=siz[x]; //累加连通块的大小
}
void solve(int u,int l,int r){
  int now=top;
  //枚举当前区间出现的颜色边，合并其端点
  for(auto v:tr[u])merge(v.first,v.second);
  //枚举叶子颜色边，累加相邻连通块大小之积
  if(l==r) for(auto v:e[r]) ans+=
    siz[find(v.first)]*siz[find(v.second)];
  else solve(ls,l,mid),solve(rs,mid+1,r);
  
  while(top>now){ //撤销合并
    auto s=st[top--];
    p[s.x]=s.x;
    high[s.y]=s.hy;
    siz[s.y]-=siz[s.x];
  }
}
signed main(){
  scanf("%lld",&n);
  for(int i=1;i<=n;i++)p[i]=i,siz[i]=1;
  for(int i=1,x,y,c;i<n;i++){
    scanf("%lld%lld%lld",&x,&y,&c);
    if(c!=1) insert(1,1,n,1,c-1,{x,y});
    if(c!=n) insert(1,1,n,c+1,n,{x,y});
    e[c].push_back({x,y}); //记录同种颜色的边
  }
  solve(1,1,n);
  cout<<ans;
}

 

// 线段树分治 O(nlognlogn)
#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;

#define int long long
#define ls (u<<1)
#define rs (u<<1|1)
#define mid (l+r>>1)
const int N=500005;

int n,m,top,ans;
int p[N],siz[N];
vector<pair<int,int> >e[N]; //存边
vector<pair<int,int> >tr[N<<2];
struct node{
  int x,y;
}st[N]; //栈

void insert(int u,int l,int r,int L,int R,pair<int,int> e){
  if(L<=l&&r<=R) return tr[u].push_back(e);
  if(L<=mid) insert(ls,l,mid,L,R,e);
  if(R>mid) insert(rs,mid+1,r,L,R,e);
}
int find(int x){ //查找根
  return p[x]==x?x:find(p[x]);
}
void merge(int x,int y){ //合并集合
  x=find(x),y=find(y);
  if(x==y) return;
  if(siz[x]>siz[y])swap(x,y);
  st[++top]={x,y};
  p[x]=y;
  siz[y]+=siz[x];
}
void solve(int u,int l,int r){
  int now=top;
  for(auto v:tr[u])merge(v.first,v.second);

  if(l==r) for(auto v:e[r]) ans+=
    siz[find(v.first)]*siz[find(v.second)];
  else solve(ls,l,mid),solve(rs,mid+1,r);
  
  while(top>now){ //撤销合并
    auto s=st[top--];
    p[s.x]=s.x;
    siz[s.y]-=siz[s.x];
  }
}
signed main(){
  scanf("%lld",&n);
  for(int i=1;i<=n;i++)p[i]=i,siz[i]=1;
  for(int i=1,x,y,c;i<n;i++){
    scanf("%lld%lld%lld",&x,&y,&c);
    if(c!=1) insert(1,1,n,1,c-1,{x,y});
    if(c!=n) insert(1,1,n,c+1,n,{x,y});
    e[c].push_back({x,y}); //记录同种颜色的边
  }
  solve(1,1,n);
  cout<<ans;
}