Acwing 240食物链(带权并查集)

AcWing 240.食物链

#include <iostream>
using namespace std;

//巧妙的利用余数的关系
//种类0 : d[x] % 3 == 0
//种类1 : d[x] % 3 == 1
//种类2 : d[x] % 3 == 2
//2 -> 1 -> 0 -> 2

const int N = 5e4 + 10;
int p[N], d[N];

int find(int x)
{
    if(x != p[x])
    {
        int u = p[x];       //父节点记得保存，做完find后p[x]就保存的是跟节点的值了,因为做了路径压缩
        p[x] = find(p[x]);  //找到自己的跟节点
        d[x] += d[u];       //将到父节点的距离更新成到跟节点的距离
    }
    return p[x];
}

int main()
{
    int n, k;
    cin >> n >> k;
    
    for(int i = 1; i <= n; ++i) p[i] = i;
    
    int ans = 0;
    while(k--)
    {
        int D, x, y;
        cin >> D >> x >> y;
        if(x > n || y > n)  ++ans;
        else
        {
            int px = find(x), py = find(y);
            if(D == 1) // x和y是同类 那么 p[x] == p[y] && (d[x] - d[y]) % 3 == 0
            {
                if(px == py)    //在同一食物链
                {
                    //(d[x] % 3 == d[y] % 3) 与 (d[x] - d[y]) % 3 == 0 等价
                    if((d[x] - d[y]) % 3)  // 当等于1或者2时说明是假话
                        ++ans; 
                }  
                else //不在同一食物链, 需要将两个连到同一集合
                {
                    p[px] = py;
                    //要满足d[x] % 3 == d[y] % 3 所以(d[x] + ? - d[y]) % 3 == 0 所以d[px] = d[y] - d[x];
                    d[px] = d[y] - d[x]; //到跟py的距离
                }
            }
            else    //x吃y 那么 (d[x] - d[y] - 1) % 3 == 0
            {
                if(px == py)
                {
                    if((d[x] - d[y] - 1) % 3)
                        ++ans;
                }
                else
                {
                    p[px] = py;
                    d[px] = 1 + d[y] - d[x];
                }
            }
        }
    }
    
    cout << ans << endl;
    
    return 0;
}