4084 号码牌（并查集）

1. 问题描述：

有 n 个小朋友，编号 1∼n。每个小朋友都拿着一个号码牌，初始时，每个小朋友拿的号码牌上的号码都等于其编号。每个小朋友都有一个幸运数字，第 i 个小朋友的幸运数字为 di。对于第 i 个小朋友，他可以向第 j 个小朋友发起交换号码牌的请求，当且仅当 |i − j| = di 成立。注意，请求一旦发出，对方无法拒绝，只能立刻进行交换。每个小朋友都可以在任意时刻发起任意多次交换请求。给定一个 1∼n 的排列 a1，a2，…，an。请问，通过小朋友相互之间交换号码牌，能否使得第 i 个小朋友拿的号码牌上的号码恰好为 ai，对 i∈[1，n] 均成立。

输入格式

第一行包含整数 n。第二行包含 n 个整数 a1，a2，…，an。第三行包含 n 个整数 d1，d2，…，dn。

输出格式

共一行，如果能做到，则输出 YES，否则输出 NO。

数据范围

前 6 个测试点满足 1 ≤ n ≤ 10。
所有测试点满足 1 ≤ n ≤ 100，1 ≤ di ≤ n，保证 a1∼an 是一个 1∼n 的排列。

输入样例1：

5
5 4 3 2 1
1 1 1 1 1

输出样例1：

YES

输入样例2：

7
4 3 5 1 2 7 6
4 6 6 1 6 6 1

输出样例2：

NO

输入样例3：

7
4 2 5 1 3 7 6
4 6 6 1 6 6 1

输出样例3：

YES
来源：https://www.acwing.com/problem/content/4087/

2. 思路分析：

分析题目可以知道每一个小朋友在满足可以交换的条件下可以进行交换的操作，并且可以发现交换具有传递性，例如A可以和B交换，B可以和C交换，那么A可以和C交换，进而A，B，C之间可以任意交换，相互交互在图论上属于一个连通块，相当于一棵树的形式，也即在同一个集合中的所有元素是可以任意交换的，不在同一个集合中的元素不能够交换，根据这个特点可以知道我们需要合并同一个集合中的所有元素，而合并集合中的元素可以使用并查集进行合并，所以这道题目考察的是并查集的相关操作，由于一开始的时候i = ai，所以我们可以从1枚举到n，将当前节点合并到它所在的集合中，声明A，B分别记录集合中已经有牌的集合和需要的牌的集合，从1枚举到n，枚举的过程中通过并查集的find操作找到当前的节点i所在的集合，当前节点i所在的集合A[find(i)]加上当前的节点i，当前节点i所在的集合B[find(i)]加上a[i]表示当前集合需要a[i]这个元素，最终判断A和B中每一个集合的元素是否相同即可。

3. 代码如下：

from typing import List


class Solution:
    def find(self, x: int, fa: List[int]):
        if x != fa[x]: fa[x] = self.find(fa[x], fa)
        return fa[x]

    def merge(self, a: int, b: int, n: int, fa: List[int]):
        # 越界了所以需要跳过
        if b < 1 or b > n: return
        # 合并a, b两个节点为同一个集合
        fa[self.find(a, fa)] = self.find(b, fa)

    # 并查集
    def process(self):
        n = int(input())
        a = [0] + list(map(int, input().split()))
        d = [0] + list(map(int, input().split()))
        fa = [i for i in range(n + 10)]
        for i in range(1, n + 1):
            # 合并属于同一个集合中的所有点
            self.merge(i, i - d[i], n, fa)
            self.merge(i, i + d[i], n, fa)
        # A为当前已有牌的集合, B为需要的牌的集合
        A, B = [list() for i in range(n + 10)], [list() for i in range(n + 10)]
        for i in range(1, n + 1):
            # 合并每一个节点属于哪一个集合
            A[self.find(i, fa)].append(i)
            B[self.find(i, fa)].append(a[i])
        # A是有序的但是B不是有序的所以需要对每一个B[i]进行排序
        for i in range(1, n + 1):
            if B[i]:
                # 从小到大进行排序
                B[i].sort()
        # 判断两个集合是否相等, python直接判断相等即可
        for i in range(1, n + 1):
            if A[i] != B[i]:
                print("NO")
                return
        print("YES")


if __name__ == "__main__":
    Solution().process()