蓝桥幼儿园

问题描述

蓝桥幼儿园的学生是如此的天真无邪，以至于对他们来说，朋友的朋友就是自己的朋友。

小明是蓝桥幼儿园的老师，这天他决定为学生们举办一个交友活动，活动规则如下：

小明会用红绳连接两名学生，被连中的两个学生将成为朋友。

小明想让所有学生都互相成为朋友，但是蓝桥幼儿园的学生实在太多了，他无法用肉眼判断某两个学生是否为朋友。于是他起来了作为编程大师的你，请你帮忙写程序判断某两个学生是否为朋友（默认自己和自己也是朋友）。

输入描述

第 1 行包含两个正整数 N,M，其中 N 表示蓝桥幼儿园的学生数量，学生的编号分别为 1∼N。

之后的第 2∼M+1 行每行输入三个整数，op,x,y：

• 如果 op=1，表示小明用红绳连接了学生 x 和学生 y 。
• 如果 op=2 请你回答小明学生 x 和 学生 y 是否为朋友。

x,y<=2*10^5

输出描述

对于每个 op=2op=2 的输入，如果 xx 和 yy 是朋友，则输出一行 YES，否则输出一行 NO。

输入输出样例

示例 1

输入

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

输出

NO
YES
YES

我的代码：

采用并查集，注意x,y的范围，当范围有问题的时候，会报段错误

N = 200005
def init_set():
    global s
    s= list(range(N))

def find_set(x):
    global s
    if x!=s[x]:
        s[x] = find_set(s[x])
    return s[x]

def merge_set(x,y):
    global s
    x = find_set(x)
    y = find_set(y)
    if x!=y:
        s[x] = s[y]

def friend(x,y):
    x = find_set(x)
    y = find_set(y)
    if x == y :
        print("YES")
    else:
        print("NO")
    

init_set()
n,m = map(int,input().split())
for k in range(m):
    op,x,y = map(int,input().split())
    if op==1:
        merge_set(x,y)
    else:
        friend(x,y)