[LeetCode] 2368. Reachable Nodes With Restrictions

- Problem

2368. Reachable Nodes With Restrictions

0번 노드에서 출발하여, 0번 노드를 포함한 도달할 수 있는 노드들의 총 개수를 구하는 문제이다.

- 내 풀이 (시간 초과)

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        def bfs(x):
            q = collections.deque([x])
            visited = [0] * n
            visited[x] = 1
            
            while q:
                x = q.popleft()
                
                for next_node in graph[x]:
                    if not visited[next_node] and next_node not in restricted:
                        q.append(next_node)
                        visited[next_node] = 1
            
            return sum(visited)
        
        graph = [[]*n for _ in range(n)]
        
        for x, y in edges:
            graph[x].append(y)
            graph[y].append(x)
             
        return bfs(0)

bfs를 사용하여 문제를 해결했으나 마지막 테스트 케이스에서 시간 초과가 뜬다.
내 생각엔 bfs를 돌때, if not visited[next_node and next_node not in restricted:에서 not in연산에 많은 시간이 소모되지 않을까 생각했다. (in연산은 O(n)의 시간 복잡도를 가진다.)

그래서 주어진 restricted 배열을 set으로 바꿔 in 연산을 진행하였다

List에서 in 연산은 O(n), Set에서 in 연산은 O(1)의 시간 복잡도를 가진다.
간단하게 말하자면 파이썬의 set은 hash table로 구현되어 있기 때문이다.
자세한 설명은 여기서 확인할 수 있다.

- 두 번째 풀이

class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        def bfs(x):
            q = collections.deque([x])
            visited = [0] * n
            visited[x] = 1
            cnt = 0
            
            while q:
                x = q.popleft()
                cnt += 1
                
                for next_node in graph[x]:
                    if not visited[next_node]:
                        q.append(next_node)
                        visited[next_node] = 1
            
            return cnt
        
        graph = [[]*n for _ in range(n)]
        restricted = set(restricted)
        
        for x, y in edges:
            if x not in restricted and y not in restricted:
                graph[x].append(y)
                graph[y].append(x)
             
        return bfs(0)

- 결과