2 minute read

Problem Statement

leetcode problem link

DFS Approach [TLE]

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        graph = {i:[] for i in range(n)}
        for src, des in edges:
            graph[src].append(des)
            graph[des].append(src)

        visited = set()
        res = [0]

        def dfs(i):
            visited.add(i)
            res[0] += 1
            for nei in graph[i]:
                if nei not in visited and nei not in restricted:
                    dfs(nei)

        dfs(0)

        return res[0]

BFS Approach [Accepted]

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        graph = {i:[] for i in range(n)}
        for src, des in edges:
            graph[src].append(des)
            graph[des].append(src)

        visited = set()
        res = 0
        resticted_nodes = set(restricted)
        if 0 in resticted_nodes:
            return 0

        queue = deque()
        queue.append(0)
        visited.add(0)
        while queue:
            node = queue.popleft()
            visited.add(node)
            res += 1
            for nei in graph[node]:
                if nei not in visited and nei not in resticted_nodes:
                    queue.append(nei)

        return res

Editorial

Approach 1: Breadth First Search (BFS)

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        # Store all edges in 'neighbors'.
        neighbors = collections.defaultdict(list)
        for node_a, node_b in edges:
            neighbors[node_a].append(node_b)
            neighbors[node_b].append(node_a)

        # Mark the nodes in 'restricted' as visited.
        seen = [False] * n
        for node in restricted:
            seen[node] = True

        # Store all the nodes to be visited in 'queue'.
        ans = 0
        queue = collections.deque([0])
        seen[0] = True

        while queue:
            curr_node = queue.popleft()
            ans += 1

            # For all the neighbors of the current node, if we haven't visit it before,
            # add it to 'queue' and mark it as visited.
            for next_node in neighbors[curr_node]:
                if not seen[next_node]:
                    seen[next_node] = True
                    queue.append(next_node)

        return ans

Approach 2: Depth First Search (DFS): Recursive

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        # Store all edges according to nodes in 'neighbors'.
        neighbors = collections.defaultdict(list)
        for node_a, node_b in edges:
            neighbors[node_a].append(node_b)
            neighbors[node_b].append(node_a)

        # Mark the nodes in 'restricted' as visited.
        seen = [False] * n
        for node in restricted:
            seen[node] = True

        def dfs(curr_node):
            # Mark 'curr_node' as visited and increment 'ans' by 1.
            self.ans += 1
            seen[curr_node] = True

            # Go for unvisited neighbors of 'currNode'.
            for next_node in neighbors[curr_node]:
                if not seen[next_node]:
                    dfs(next_node)

        self.ans = 0
        dfs(0)
        return self.ans

Approach 3: Depth First Search (DFS): Iterative

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        # Store all edges according to nodes in 'neighbor'.
        neighbors = collections.defaultdict(set)
        for a, b in edges:
            neighbors[a].add(b)
            neighbors[b].add(a)

        # Mark the nodes in 'restricted' as visited.
        seen = [False] * n
        for node in restricted:
            seen[node] = True

        # Use stack 'stack' to store all nodes to be visited, start from node 0.
        stack = [0]
        ans = 0
        seen[0] = True

        while stack:
            curr_node = stack.pop()
            ans += 1

            # Add all unvisited neighbors of the current node to 'stack'
            # and mark them as visited.
            for next_node in neighbors[curr_node]:
                if not seen[next_node]:
                    seen[next_node] = True
                    stack.append(next_node)

        return ans

Approach 4: Disjoint Set Union (DSU)

class UnionFind:
    def __init__(self, n):
        self.root = list(range(n))
        self.rank = [1] * n
    def find(self, x):
        if self.root[x] != x:
            self.root[x] = self.find(self.root[x])
        return self.root[x]
    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x != root_y:
            if self.rank[root_x] > self.rank[root_y]:
                root_x, root_y = root_y, root_x
            self.rank[root_y] += self.rank[root_x]
            self.root[root_x] = root_y
    def getSize(self, x):
        return self.rank[self.find(x)]

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        rest_set_ = set(restricted)
        uf = UnionFind(n)

        for a, b in edges:
            if a not in rest_set and b not in rest_set:
                uf.union(a, b)

        return uf.getSize(0)