[백준] 11724-연결요소의 개수

문제바로가기

n, m = map(int, input().split())
count = 0    
data = [[] for _ in range(n + 1)]
visited = [[0 for _ in range(n+1)] for _ in range(n+1)]

def dfs(data, idx):
    next = data[idx][1]
    visited[data[idx][0]][data[idx][1]] = True
    
    id = list(filter(lambda x: data[x][0] == next, range(len(data))))
    for d in id:
       if visited[data[d][0]][data[d][1]] is not True:
            dfs(data, d)

for _ in range(m):
    data.append(list(map(int, input().split())))
    
for i, num in enumerate(data):
    if visited[num[0]][num[1]] == 0:
        dfs(data, i)
        count += 1
print(count)

위 처럼 idx 기반으로 방문을 표시하는 방법도 정답이 나온다. 하지만 큰 리스트 및 filter 연산으로 메모리 초과 혹은 시간초과가 나온다. (각각 pypy3, python3으로 설정 시, 발생) 따라서 방문을 2차원 리스트의 인덱스를 일일이 검사하는 방식을 바꿔야 한다.

import sys
sys.setrecursionlimit(10**5)

n, m = map(int, input().split())
data = [[] for _ in range(n+1)]
visited = [False]*(n+1)
count = 0

for _ in range(m):
    u, v = map(int, input().split())
    data[u].append(v)
    data[v].append(u)
    
def dfs(idx):
    visited[idx] = True
    for i in data[idx]:
        if not visited[i]:
            dfs(i)


for i in range(1, n + 1):
    if not visited[i]:
        dfs(i)
        count += 1

print(count)

맨 위의 import sys sys.set...~ 부분은 dfs의 무한 메모리 사용을 방지하기 위한 상한을 만들어둔 것으로 보면된다. 이 부분은 예전에 다른 문제를 풀 때 본 블로그에서 알게 된 방식으로 요긴하게 잘 쓰는 중.

위 코드랑 다른 점은 일단 visited이다. 일단 무방향(혹은 양방향)이기 때문에 data에 u에서 v, v에서 u 모두 업데이트 해주어야 한다. 또한, 예시로 visited[1]이라고 하면 1과 연관된 모든 간선을 의미하기 때문에 앞선 코드와 달리 많은 검사를 할 필요가 없다.
이렇게 하면 훨씬 더 효율적인 코드를 작성가능!