带权并查集 记录节点到根的距离

AcWing 240. 食物链

#include <bits/stdc++.h>
using namespace std;

const int N = 50005;

int p[N], dis[N];   //dis为到根节点的距离

int find(int x) {
    if (x != p[x]) {
        int former_p = p[x];
        p[x] = find(p[x]);
        dis[x] += dis[former_p];    //更新到根节点的距离
    }

    return p[x];
}

int main() 
{
#ifndef ONLINE_JUDGE
    freopen("D:/VS CODE/C++/in.txt", "r", stdin);
    freopen("D:/VS CODE/C++/out.txt", "w", stdout);
#endif
    int n, k;
    int ans = 0;
    cin >> n >> k;

    for (int i = 1; i <= n; ++i) {
        p[i] = i;
    }

    while (k -- ) {
        int c, x, y;
        cin >> c >> x >> y;

        if (x > n or y > n) ++ans;
        else if (c == 1) {
            //x y 同类
            int r_x = find(x), r_y = find(y);
            if (r_x != r_y) {
                //x 和 y之间的关系还不明确
                p[r_x] = r_y;
                dis[r_x] = dis[y] - dis[x]; 
                //(dis[r_x] + dis[x] - dis[y]) % 3 = 0;
            }
            else {
                if ((dis[x] - dis[y]) % 3) {
                    //不为同类
                    ++ans;
                }
            }
        }
        else {
            int r_x = find(x), r_y = find(y);
            if (r_x != r_y) {
                p[r_x] = r_y;
                dis[r_x] = dis[y] - dis[x] + 1;
            }
            else {
                if ((dis[x] - dis[y] - 1) % 3)
                    ++ans; 
            }
        }
    }

    printf("%d", ans);

    fclose(stdin);
    fclose(stdout);
    return 0;
}