[백준 1976] 여행 가자

문제 출처 : https://www.acmicpc.net/problem/1976

1976번: 여행 가자

 

여행을 갈 수 있는지, 없는지 확인하는 문제 입니다. 문제를 보면 쉽게 도시들을 DFS를 통해 방문이 가능한지 불가능한지 따지면 해결될 것으로 보입니다. 하지만 이렇게 문제를 풀게 되면 시간초과가 발생할 수 있습니다. 왜냐하면 의도적으로 끝에서 끝으로 이동하는 경로만 입력이 들어온다면 방문하는데 시간이 오래걸릴 수 있습니다.

따라서 이 문제는 DFS가 아니라 유니온 파인드, 즉 분리 집합으로 문제를 해결해야 합니다. 그래도 되는 이유는 어떻게 이동해야 하는지 경로에 대해 묻는 부분이 없습니다. 결국 연결이 되어 있는지, 연결이 되어 있지 않은지만 파악하면 되기 때문 입니다.

그럼 문제를 풀어보도록 하겠습니다.

입력 받기

mii = lambda : tuple(map(int, input().split()))

N = int(input())
M = int(input())

for i in range(N):
    row = mii()
    for j, c in enumerate(row):
        if c:
            union(i, j)

q = mii()

여러 입력이 들어왔을 때를 처리하기 위한 map(int, input().split())을 여러 번 쓰기 싫어 lambda 함수를 사용하여 함수화 해주었습니다. lambda의 활용에 대해서 잘 모른다면 여기 링크를 확인해 보시기 바랍니다.

따로 값이 변경되지 않기 때문에 tuple로 만들어 속도를 조금이라도 빠르게 해주었습니다. 만약 이 값들이 변경된다면 list로 만들어 주어야 합니다.

전체 도시의 개수 N과 여행할 도시 M을 입력 받습니다. 다음으로 인접 형렬 형태의 도시 N줄을 입력 받습니다. 각각의 줄은 row로 입력을 받아 해당 row가 1이면 두 도시가 연결되어 있기 때문에 union을 해줍니다.

다음으로 M개의 도시를 여행하는 목록 q를 입력 받습니다.

Union Find

다음으로 유니온 파인드의 로직을 만들어 줍니다.

parent = [i for i in range(N+1)]

def find(a):
    if a == parent[a]:
        return a
    parent[a] = find(parent[a])
    return parent[a]

def union(i, j):
    i = find(i)
    j = find(j)
    parent[j] = i

먼저 parent는 도시들을 뜻합니다. 처음에는 자기 자신만 연결되어 있기 때문에 for문으로 0부터 N-1까지의 배열로 만들어 줍니다. find를 통해 연결된 부모를 찾아 리턴해 줍니다. 그리고 union을 통해 두 도시를 연결해 줍니다.

연결 여부 확인하기

p = find(q[0] - 1)
ans = "YES"

for i in range(1, M):
    if p != find(q[i] - 1):
        ans = "NO"
        break
  
print(ans)

입력받은 q의 첫 번째 도시부터 부모를 찾습니다. 저장된 도시들은 0부터 시작하는데 입력은 1부터 들어오기 때문에 q의 값에 1을 빼 주었습니다. 결과를 출력하는 ans 값의 초기값으로는 YES를 저장합니다. 만약 연결된 도시가 없다면 NO가 될 것입니다.

0 번째 도시의 부모를 찾고 난 뒤 다음 도시부터는 for문을 돌면서 부모가 같은지 확인합니다. 만약 하나라도 같은 부모가 아닌 도시가 나온다면 ans값이 “NO”가 됩니다. 최종적으로 모든 도시가 연결되어 있다면 YES를, 하나라도 연결되지 않은 도시가 있다면 NO를 출력하고 프로그램이 종료될 것입니다.

전체 코드

전체 코드를 확인해 보겠습니다.

mii = lambda : tuple(map(int, input().split()))

N = int(input())
M = int(input())

parent = [i for i in range(N+1)]

def find(a):
    if a == parent[a]:
        return a
    parent[a] = find(parent[a])
    return parent[a]

def union(i, j):
    i = find(i)
    j = find(j)
    parent[j] = i

for i in range(N):
    row = mii()
    for j, c in enumerate(row):
        if c:
            union(i, j)

q = mii()
p = find(q[0] - 1)

ans = "YES"
for i in range(1, M):
    if p != find(q[i] - 1):
        ans = "NO"
        break
  
print(ans)