C++算法题 & 食物链 & 带权并查集

题目

题目来源AcWing第240题

动物王国中有三类动物 A,B,C，这三类动物的食物链构成了有趣的环形。
A 吃 B，B 吃 C，C 吃 A。
现有 N 个动物，以 1∼N 编号。
每个动物都是 A,B,C 中的一种，但是我们并不知道它到底是哪一种。
有人用两种说法对这 N 个动物所构成的食物链关系进行描述：

• 第一种说法是 1 X Y，表示 X 和 Y 是同类。
• 第二种说法是 2 X Y，表示 X 吃 Y。

此人对 N 个动物，用上述两种说法，一句接一句地说出 K 句话，这 K 句话有的是真的，有的是假的。
当一句话满足下列三条之一时，这句话就是假话，否则就是真话。

• 当前的话与前面的某些真的话冲突，就是假话；
• 当前的话中 X 或 Y 比 N 大，就是假话；
• 当前的话表示 X 吃 X，就是假话。

你的任务是根据给定的 N 和 K 句话，输出假话的总数。

输入格式

第一行是两个整数 N 和 K，以一个空格分隔。
以下 K 行每行是三个正整数 D，X，Y，两数之间用一个空格隔开，其中 D 表示说法的种类。

• 若 D=1，则表示 X 和 Y 是同类。
• 若 D=2，则表示 X 吃 Y。

输出格式

只有一个整数，表示假话的数目。

数据范围

1 ≤ N ≤ 50000,
0 ≤ K ≤ 100000

输入样例：

100 7
1 101 1
2 1 2
2 2 3
2 3 3
1 1 3
2 3 1
1 5 5

输出样例：

3

解题思路

分析题意，我们可以并查集的方法来解决这个问题，每个动物个体都用集合来表示。因为每个节点都可以找到它相应的根节点，所以我们可以将已知关系的集合放到一个集合中去。节点的相互关系我们可以通过定义节点到根节点距离来表示，因为动物有三个种类，所以节点与根节点的关系有三种(da表示a到根节点的距离，db表示b到根节点的距离)：

1. (da - db) % 3 == 0 a,b为同类
2. (da - db) % 3 == 1 b吃a
3. (da - db) % 3 == 2 a吃b

后面就是题目要求的操作：
首先是 1 X Y，表示 X 和 Y 是同类操作，如果x，y关系还不知道（前面没有说），则这句话是真话且x，y不在同一集合，那我们只需将x，y连接到一起且距离d[px] = d[y] - d[x]，这样就让x,y到根节点的距离相等，也就是让x，y为同类。如果x，y关系前面有确定过，则x，y在同一集合，我们只需判断x，y到根结点的距离差为3的倍数 [(dx - dy)%3 == 0] 则x，y为同类，否则这是假话。
接下来是2 X Y，表示 X 吃 Y，和1操作类似，不同点是在处理x，y的距离差是 模三加1[(dx - dy - 1)%3 == 0]

代码

#include <iostream>
using namespace std;

const int N = 1000010;
// p[N]存放节点所在trie树的根节点，d[N]存放点到根节点的距离
int p[N], d[N];

int find(int x){
    if(p[x] != x){
        int u = find(p[x]);
        d[x] += d[p[x]];
        p[x] = u; 
    }
    return p[x];
}
int main(){
    int n, m;
    scanf("%d%d",&n,&m);
    for(int i = 1; i <= n; i++) p[i] = i;
    
    int res = 0;
    
    while(m -- ){
        int op,a,b;
        scanf("%d%d%d", &op,&a,&b);
        if(a > n || b >n) res ++;
        else{
            int pa = find(a), pb = find(b);
            if(op == 1){
                //a,b在同一trie树上,到根节点的距离差应该为3的倍数
                if((pa == pb) && (d[a] - d[b]) % 3) res++;
                else if(pa != pb){ //a,b不在同一trie树上
                    p[pa] = pb; //将a所在树连接到b上
                    d[pa] = d[b] - d[a]; //a,b为同类，所以到根节点pb距离要相等(%3后相等)
                }
            }
            else if(op == 2){//a吃b
                if(pa == pb && (d[a] - d[b] - 1)%3) res++;
                else if(pa != pb){ //a,b不在同一trie树上
                    p[pa] = pb; //将a所在树连接到b上
                    d[pa] = d[b] - d[a] + 1; 
                }
            }
        }
    }
    printf("%d\n", res);
    return 0;
}