AcWing 240. 食物链

带权并查集，用权值维护是不是属于同一类

#include <iostream>
#include <algorithm>
#include <cstring>
using namespace std;
const int N = 50010;
int fa[N],d[N];//d[i]维护i到自己祖先节点的距离
int n,k,res;
int get(int x) {
    if(x==fa[x]) return x;
    int t=get(fa[x]);
    d[x]+=d[fa[x]];
    fa[x]=t;
    return fa[x];
}
int main() {
    #ifdef ONLINE_JUDGE
    #else
    freopen("in.txt","r",stdin);
    freopen("out.txt","w",stdout);
    #endif
    scanf("%d%d",&n,&k);
    for(int i=1;i<=n;++i) fa[i]=i;
    for(int i=1;i<=k;++i) {
        int op,x,y;
        scanf("%d%d%d",&op,&x,&y);
        if(x>n||y>n) {
            ++res;
            continue;
        }
        if(op==1) {
            int fx=get(x),fy=get(y);
            //不在同一颗树上
            if(fx!=fy) {
                //合并到同一颗树上，并且属于同类 (d[x]+d[fx]-d[y])%3==0
                fa[fx]=fy;
                d[fx]=d[y]-d[x];
            } else {
                //若再同一颗树上判断是否是同一类
                if((d[x]-d[y])%3!=0) ++res;//非同一类
                //if((d[x]%3+3)%3!=(d[y]%3+3)%3) ++res;
            }
        } else if(op==2) {
            int fx=get(x),fy=get(y);
            //在不在同一颗树上
            if(fx!=fy) {
                //合并到同一棵树上 并且让x吃y (d[x]+d[fx]-1-d[y])%3==0
                fa[fx]=fy;
                d[fx]=d[y]+1-d[x];
            } else {
                //若在同一颗树上，判断是否为x吃y
                if((d[x]-1-d[y])%3!=0) ++res;
                //也可写成 if(((d[x]-1)%3+3)%3!=(d[y]%3+3)%3) ++res;
            }
        }
    }
    printf("%d",res);
    return 0;
}