ARC 097 D Equals （python）

atcoder.jp

けんちょんさんのガイドに沿って解き進めています。
union find の問題。
交換する数字そのものが与えられるので、どことどこを交換したらいいかを効率的に調べる必要があります。
まず、数字から位置を求める配列を作りました。

例えば、上の様な配列があるとしたら、２は先頭にあるので、list[2]=1となります。

n, m = map(int, input().split())
pn = list(map(lambda x:int(x)-1, input().split()))
ls = [-1] * n
for i in pn:
    ls[pn[i]] = i

そこからしばらくは、union findの定型です。

def find(x):
    if par[x] == x:
        return x
    else:
        s = find(par[x])
        par[x] = s
        return s

def unite(x,y):
    s = find(x)
    t = find(y)
    if s>t:
        par[s] = t
    else:
        par[t] = s


for _ in range(m):
    x, y = map(lambda x:int(x)-1, input().split())
    unite(ls[x],ls[y]) #位置を基準にunite処理をしていきます。

最終的には、i番目の場所(place1)の親 と数字iの場所(place2)の親がつながっていれば良いことになります。

ans2 = 0
for i in range(n): 
    place1 = i
    place2 = ls[i]

    if find(place1)==find(place2):
       ans2+=1

print(ans2)