[Python/백준] 1976. 여행 가자

문제 접근 및 공부 내용, 풀이는 모두 하단 코드에 "주석"으로 포함되어 있으니 참고해주세요.

문제 유형 보기

유니온 파인드

https://www.acmicpc.net/problem/1976

---

풀이

성공, 주석 제거 버전

import sys
input = sys.stdin.readline

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])

    return parent[x]


def union(x, y):
    x_root = find(x)
    y_root = find(y)

    if x_root == y_root:
        return

    if rank[x_root] > rank[y_root]:
        parent[y_root] = x_root
    elif rank[x_root] < rank[y_root]:
        parent[x_root] = y_root
    else:
        parent[y_root] = x_root
        rank[x_root] += 1

n = int(input())
m = int(input())

graph = [list(map(int, input().split())) for _ in range(n)]
travels = list(map(int, input().split()))

parent = list(range(n+1))
rank = [0] * (n+1)

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

travel_set = {find(x) for x in travels}

print("YES" if len(travel_set) == 1 else "NO")

성공, 주석 포함 버전

import sys
input = sys.stdin.readline

def find(x):
    if parent[x] != x:
        # 트리 구조 변경하면서 root를 찾아옴
        parent[x] = find(parent[x])

    return parent[x]


def union(x, y):
    x_root = find(x)
    y_root = find(y)

    if x_root == y_root:
        return

    if rank[x_root] > rank[y_root]:
        parent[y_root] = x_root
    elif rank[x_root] < rank[y_root]:
        parent[x_root] = y_root
    else:
        parent[y_root] = x_root
        rank[x_root] += 1
# 여기까지는 정말 기본적인 유니온 파인드 ADT 구현


# 여행을 가는 데,
# 입력 :
# 첫 줄 : 도시 수 N
# 둘째 줄 : 여행 계획에 속한 도시들 수 M
# 그 담부터 : 인접 행렬 N x N
# 마지막 : 여행 계획(방문할 도시들)
# 도시 번호는 1, 2, ..., n임

n = int(input())
m = int(input())

graph = [list(map(int, input().split())) for _ in range(n)]
travels = list(map(int, input().split()))

# 아, 이거 travels에 속한 애들이, 지금은 서로소 집합인데,
# 얘네들이 결국 하나의 집합을 이루면 YES 못하면 NO 출력(삼항)
# 그러면, parents 놓고, 거기에 모든 노드들 넣고, 연결시킨다음에,
# 마지막에 최종 parents 보고서 거기 중에 travels에 속한 놈들의 find(travel)이
# 모두 같은 지 비교하면 됨.

# 모두 같은 지 비교하는 법은 여러 개 있을듯?
# 그냥 하나씩 find 해서 그 결과 반복문으로 비교해, 이전과 다르면 NO 출력
# 다른 거 없이 반복문 다 헤쳐나갔으면 YES 출력 이렇게 할 수도 있을거고
# 아님 set() 한 다음에 길이 재서 1이냐 아니냐 할 수도 있을듯.

parent = list(range(n+1))
rank = [0] * (n+1)

# 반복문 돌면서 union 호출하면 될듯.
for i in range(n):
    for j in range(i):
        if graph[i][j] == 1:
            union(i + 1, j + 1)

travel_set = {find(x) for x in travels}

print("YES" if len(travel_set) == 1 else "NO")

강사님 정답 풀이 확인

group = find(plans[0])  # 출발 도시의 집합이 기준이 됨
result = "YES"

for plan in plans[1:]:
    if group != find(plan):
        result = "NO"  # 하나의 도시라도 다른 집합에 속하면 여행할 수 없음
        break

print(result)

# 이렇게 비교할 수도 있다.
# 사실, 엄밀히 따지면 내 풀이보다 이렇게 푸는 게 더 실행 시간 상에서 이득임.
# 왜냐하면, 얘는 마치 함수의 early return처럼, 여행이 불가능하면 바로 break 해버리기에,
# len(travels)만큼 모두 살피지 않아도 될 수 있으니 더 효율적임.
# 그러나, 내 코드는 set()을 쓰니 len(travels) 만큼을 다 뒤져봐야함.

# 어쨌든간에, 결국 웬만해서는 저 메인 이중 루프가 대충 (N^2)/2 정도 소요될 텐데,
# M은 1000 이하니 M만큼 더 걸려봤자 그다지 크지도 않고 (N^2)/2의 최악의 경우보다도 작음.

# 근데 굳이 set을 쓸 이유도 없고, 
# 나중이 어렇게 끝처리 해야하는 문제 또 만나면 그냥 위처럼 할듯?
# 모두 값이 같은지 체크는 set()을 사용하면 O(n)으로 검사가 가능하다는 것만 알고 가자.

---

250228 재풀이

그냥 백준 스타일대로 바로 입력받았음

import sys
input = sys.stdin.readline

# 임의의 두 도시 사이에 길이 있을 수도 없을 수도 있음.
# 즉, 여행 계획이 불가능한 계획일 수도 있음.
# 여행 가능한지 불가능한지 출력
# 도시 수 n, 여행 계획 속한 도시 수 m,
# 인접 행렬로 연결 정보 나오고,
# 마지막으로 여행 계획 주어짐.

# 간선 정보는 있는데, 무방향이고, 가중치 없음.
# 여행 가능하단 건, E C B C D 얘네가 모두 하나의 집합에 속해 있으면 되는 거임.
# 즉, 간선 정보를 가지고 union을 다 한 다음에, 쟤네들 find 다 때려서
# 모두 같으면 YES, 하나라도 다르면 NO 출력하면 됨.
# 근데 여행 일정은 순서대로 가야 하는데, 갔던 길 다시 갈 수도 있고,
# 뭐 결국 같은 집합이기만 하면, 여러 번 방문 가능하니까 노상관 같은데?

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(x, y):
    x_root = find(x)
    y_root = find(y)

    if x_root == y_root:
        return

    if rank[x_root] > rank[y_root]:
        parent[y_root] = x_root
    elif rank[y_root] > rank[x_root]:
        parent[x_root] = y_root
    else:
        parent[y_root] = x_root
        rank[x_root] += 1

n = int(input())
m = int(input())
parent = list(range(n + 1))
rank = [0] * (n + 1)

for i in range(n):
    line = list(map(int, input().split()))
    for j in range(n):
        if line[j] == 1:
            union(i + 1, j + 1)

plans = list(map(int, input().split()))

group = find(plans[0])
for plan in plans[1:]:
    if find(plan) != group:
        print('NO')
        break
else:
    print('YES')