POJ - 1182 - 食物链（带权并查集）

题目链接_POJ-1182-食物链
不能用多组数据输入，不然就WA

#include <iostream>
#include <stdio.h>
#include <string.h>
using namespace std;
const int maxn = 5e4 + 10;

int n,m,p,x,y;
int fa[maxn], r[maxn];

int find_fa(int x){
    if(fa[x] == x) return x;
    int fx = find_fa(fa[x]);
    // 路径压缩
    // 0 - 同类, x -> 1 -> y (x 吃 y), x -> 2 -> y (x 被 y 吃)
    // 1 由输入产生， 2 由路径压缩产生 一共只用三种关系 A -> B -> C -> A
    // x -> 0 -> fx -> 0 -> fxx   =>  x -> 0 -> fxx
    // x -> 1 -> fx -> 0 -> fxx   =>  x -> 1 -> fxx
    // x -> 2 -> fx -> 0 -> fxx   =>  x -> 2 -> fxx

    // x -> 0 -> fx -> 1 -> fxx   =>  x -> 1 -> fxx
    // x -> 1 -> fx -> 1 -> fxx   =>  x -> 2 -> fxx
    // x -> 2 -> fx -> 1 -> fxx   =>  x -> 0 -> fxx

    // x -> 0 -> fx -> 2 -> fxx   =>  x -> 2 -> fxx
    // x -> 1 -> fx -> 2 -> fxx   =>  x -> 0 -> fxx
    // x -> 2 -> fx -> 2 -> fxx   =>  x -> 1 -> fxx

    r[x] = (r[fa[x]] + r[x]) % 3;

    return fa[x] = fx;
}

int main()
{
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++) {
        fa[i] = i; r[i] = 0;
    }
    int ans = 0;
    for(int i=1;i<=m;i++){
        scanf("%d%d%d",&p,&x,&y);
        if(x>n || y>n) {
            ans ++;
            continue;
        }
        int fx = find_fa(x), fy = find_fa(y);
        if(fx != fy){
            r[fx] = (3-r[x]+p-1+r[y]) % 3;
            fa[fx] = fy;
            continue;
        }
        if(p == 1){
            // 合并
            // 当x与y与根结点的关系相同时说明它们已经是同一类了无需再合并
            if(r[x] != r[y]) ans ++; // 不相同则一定有吃或被吃关系

        } else {
            // x 吃 y
            // y 与 fa 关系为 r[y], fa 与 y 关系为 (3-r[y]) % 3
            // 则 x 与 y 关系则为 (r[x] + (3-r[y]) % 3) % 3
            if((r[x] + 3 - r[y]) % 3 != 1) ans ++;

        }
    }

    printf("%d\n",ans);

    return 0;
}