Problem p03356 (python)

Fixed submission:
import sys
from collections import defaultdict, deque
import math

# import copy
from bisect import bisect_left, bisect_right
# import heapq

# sys.setrecursionlimit(1000000)

# input aliases
input = sys.stdin.readline

getS = lambda: input().strip()
getN = lambda: int(input())
getList = lambda: list(map(int, input().split()))
getZList = lambda: [int(x) - 1 for x in input().split()]

INF = 10 ** 20
MOD = 10**9 + 7
divide = lambda x: pow(x, MOD-2, MOD)

class UnionFind:
    def __init__(self, n):
        self.par = [i for i in range(n+1)]
        self.rank = [0] * (n+1)

    # 検索
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # 併合
    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # 同じ集合に属するか判定
    def same_check(self, x, y):
        return self.find(x) == self.find(y)

    def update(self):
        for i in range(len(self.par)):
            self.par[i] = self.find(i)

def solve():
    n,m = getList()
    nums = getList()
    uf = UnionFind(n)
    for i in range(m):
        x, y = getList()
        uf.union(x, y)
    uf.update()
    st1 = [[] for i in range(n + 1)]
    st2 = [[] for i in range(n + 1)]
    for i in range(n):
        st1[uf.par[i+1]].append(i+1)
        st2[uf.par[i+1]].append(nums[i])

    ans = 0
    for s1, s2 in zip(st1, st2):
        ans += len(list(set(s1) & set(s2)))



    print(ans)
def main():
    n = getN()
    for _ in range(n):
        solve()


if __name__ == "__main__":
    # main()
    solve()

Buggy submission (verdict: Runtime Error):
def getN():
    return int(input())


def getNM():
    return map(int, input().split())


def getList():
    return list(map(int, input().split()))

class UnionFind():
    def __init__(self, N):
        self.par = [i for i in range(N+20)]

    def root(self, n):
        if self.par[n] == n:
            return n
        else:
            return self.root(self.par[n])

    def isgroup(self, a, b):
        if self.root(a) == self.root(b):
            return True
        else:
            return False

    def merge(self, m, n):
        rm = self.root(m)
        rn = self.root(n)
        if rm == rn:
            pass
        else:
            self.par[rn] = rm

N,M = getNM()
pos = getList()
uf = UnionFind(N)
for i in range(M):
    x, y = getNM()
    uf.merge(x,y)

ans = 0

for i, p in enumerate(pos):
    if uf.isgroup(i+1, p):
        #print(i+1, p)
        ans += 1
    else:
        pass

#roots = [uf.root(i) for i in range(N+1)]
print(ans)

Buggy submission (verdict: Runtime Error):
def getN():
    return int(input())


def getNM():
    return map(int, input().split())


def getList():
    return list(map(int, input().split()))

class UnionFind():
    def __init__(self, N):
        self.par = [i for i in range(N+2)]

    def root(self, n):
        if self.par[n] == n:
            return n
        else:
            return self.root(self.par[n])

    def isgroup(self, a, b):
        if self.root(a) == self.root(b):
            return True
        else:
            return False

    def merge(self, m, n):
        if self.isgroup(m,n):
            pass
        else:
            self.par[self.root(n)] = self.root(m)

N,M = getNM()
pos = getList()
uf = UnionFind(N)
for i in range(M):
    x, y = getNM()
    uf.merge(x,y)

ans = 0

for i, p in enumerate(pos):
    if uf.isgroup(i+1, p):
        #print(i+1, p)
        ans += 1
    else:
        pass

#roots = [uf.root(i) for i in range(N+1)]
print(ans)