美团240309春招实习笔试——小美的朋友关系（反向并查集）

题目描述

小美认为，在人际交往中，随着时间的流逝，朋友的关系也会慢慢变淡，最终朋友关系就会淡忘。现在初始有一些朋友关系，存在一些事件会导致两个人淡忘了他们的朋友关系。小美想知道某一时刻中，某两人是否可以通过朋友介绍互相认识。

事件共有 2 种：

1 u v：代表编号 u 的人和编号 v 的人淡忘了他们的朋友关系。
2 u v：代表小美查询编号 u 的人和编号 v 的人是否能通过朋友介绍互相认识。
注：介绍可以有多层，比如 2 号把 1 号介绍给 3 号，然后 3 号再把 1 号介绍给 4 号，这样 1 号和 4 号就认识了。

输入描述

第一行输入三个正整数 n, m, q，代表总人数，初始的朋友关系数量，发生的事件数量。

接下来的 m 行，每行输入两个正整数 u, v，代表初始编号 u 的人和编号 v 的人是朋友关系。

接下来的 q 行，每行输入三个正整数 op, u, v，含义如题目描述所述。

约束条件

$$\begin{gathered} 1\le n\le 10^9 \\ 1\le m,q\le 10^5 \\ 1\le u,v\le n \\ 1\le op\le 2 \end{gathered}$$
保证至少存在一次查询操作。

输出描述

对于每次 2 号操作，输出一行字符串代表查询的答案。如果编号 u 的人和编号 v 的人能通过朋友介绍互相认识，则输出”Yes”。否则输出”No”。

示例 1

输入

5 3 5
1 2
2 3
4 5
1 1 5
2 1 3
2 1 4
1 1 2
2 1 3

输出

Yes
No
No

说明

第一次事件，1 号和 5 号本来就不是朋友，所以无事发生。

第二次事件是询问，1 号和 3 号可以通过 2 号的介绍认识。

第三次事件是询问，显然 1 号和 4 号无法互相认识。

第四次事件，1 号和 2 号淡忘了。

第五次事件，此时 1 号无法再经过 2 号和 3 号互相认识了。

解题

使用反向并查集
由于并查集仅支持插入关系而不能删除已有的关系，因此反向思考。
先遍历所有关系和事件，提取出所有事件结束后仍保持的关系，将它们加入并查集中。
然后逆序遍历事件，正序下的“淡忘”相当于逆序下的新加关系。如此即可得到结果。

n, m, q = 5, 3, 5
initial_relationship = [
    [1, 2],
    [2, 3],
    [4, 5]
]
events = [
    [1, 1, 5],
    [2, 1, 3],
    [2, 1, 4],
    [1, 1, 2],
    [2, 1, 3]
]

class UnionFind:
    def __init__(self, n: int) -> None:
        self.parent = [i for i in range(n)]  # 位置i的元素记录其父节点下标
    
    # 查找根节点
    def find(self, i: int) -> int:
        if self.parent[i] == i: return i
        # 递归查找父节点
        self.parent[i] = self.find(self.parent[i])  # 完全压缩
        return self.parent[i]
    
    def merge(self, i: int, j: int):
        root_i = self.find(i)
        root_j = self.find(j)

        self.parent[root_i] = root_j

    def query(self, i: int, j: int) -> bool:
        return self.find(i) == self.find(j)


unionFind = UnionFind(n+1)
relationship = set()
for relation in initial_relationship:
    relationship.add(tuple(sorted(relation)))  # 排序元素，保证关系不重复
for event in events:
    if event[0] == 1:
        relationship.discard(tuple(sorted(event[1:])))

# 将所有未被淡忘的关系加入并查集
for relation in relationship:
    unionFind.merge(*relation)

# 倒推事件，出现淡忘关系，将该关系加入并查集
ans = []
for event in events[::-1]:
    if event[0] == 1:
        unionFind.merge(event[1], event[2])
    else:
        ans.append('Yes' if unionFind.query(event[1], event[2]) else 'No')

# 逆序输出结果
print(ans[::-1])

参考文章