LeetCode 547. Friend Circles [Medium]

https://leetcode.com/problems/friend-circles/

并查集的题做得比较少。练练手。

class Solution:
    def findCircleNum(self, M: List[List[int]]) -> int:
        """
        Union find
        """

        l = len(M)
        parents = list(range(l))

        def get_parent(x):
            while x != parents[x]:
                x, parents[x] = parents[x], parents[parents[x]]

            return x

        def union(x, y):
            px = get_parent(x)
            py = get_parent(y)

            if px != py:
                parents[px] = py

        for i in range(l):
            for j in range(i+1, l):
                if M[i][j] == 1:
                    union(i, j)

        res = set()

        for i in range(l):
            res.add(get_parent(i))

        return len(res)

也可以用 DFS 来做：

class Solution:
    def findCircleNum(self, M: List[List[int]]) -> int:
        n = len(M)
        visited = dict()        
        circles = 0

        for i in range(n):
            if visited.get(i):
                continue

            stack = [(i, i)]
            while(len(stack) > 0):
                i, j = stack.pop()
                for index in range(n):
                    if visited.get(index): continue
                    if M[i][index]: stack.append((i, index))  
                    if M[index][j]: stack.append((j, index))   
                visited[i] = True
                visited[j] = True

            circles += 1

        return circles