백준 / 여행 가자 / 1976번 / Python

*문제 출처는 백준에 있습니다.

문제 제목: 여행 가자 / 1976번 (골드 4단계)

문제 사이트: https://www.acmicpc.net/problem/1976

 

문제 설명

---

나의 풀이

from collections import deque


def renewal(citys, n, m, planes):
    # BFS로 모든 도시가 연결되어 있는지 확인
    start = planes[0]
    for i in range(1, m):
        end = planes[i]
        if not search(citys, start, end, n):
            return False
        start = end
    return True


def search(citys, start, end, n):
    # BFS로 start에서 end로 갈 수 있는지 확인
    dq = deque([start])
    visit = [False] * (n + 1)
    visit[start] = True

    while dq:
        current = dq.popleft()
        if current == end:
            return True

        for neighbor in citys[current]:
            if not visit[neighbor]:
                visit[neighbor] = True
                dq.append(neighbor)

    return False


def main():
    # 나라의 개수
    n = int(input())
    # 여행 계획에 속한 도시들의 수
    m = int(input())

    # 도시 연결 그래프 초기화
    city = [[] for _ in range(n + 1)]

    for i in range(1, n + 1):
        edge = list(map(int, input().split()))
        for j in range(len(edge)):
            if edge[j] == 1:
                city[i].append(j + 1)
                city[j + 1].append(i)  # 무향 그래프 처리

    # 여행 계획
    planes = list(map(int, input().split()))

    if renewal(city, n, m, planes):
        print('YES')
    else:
        print('NO')


main()

위 코드도 정답이지만 개선사항으로는 Union-Fine(Disjoint Set Union) 자료구조를 사용하여 모든 도시가 같은 연결된 컴포넌트에 속해 있는지 확인하는 방식으로 접근하는 것이 적합합니다.

수정된 코드

def find(parent, x):
    if parent[x] != x:
        parent[x] = find(parent, parent[x])  # 경로 압축
    return parent[x]


def union(parent, rank, x, y):
    rootX = find(parent, x)
    rootY = find(parent, y)

    if rootX != rootY:
        if rank[rootX] > rank[rootY]:
            parent[rootY] = rootX
        elif rank[rootX] < rank[rootY]:
            parent[rootX] = rootY
        else:
            parent[rootY] = rootX
            rank[rootX] += 1


def main():
    # 입력
    n = int(input())  # 도시의 수
    m = int(input())  # 여행 계획에 포함된 도시의 수

    # 초기화
    parent = list(range(n + 1))
    rank = [0] * (n + 1)

    # 그래프 정보 입력 및 Union-Find로 연결
    for i in range(1, n + 1):
        edges = list(map(int, input().split()))
        for j in range(1, n + 1):
            if edges[j - 1] == 1:  # 연결된 경우
                union(parent, rank, i, j)

    # 여행 계획 입력
    plan = list(map(int, input().split()))

    # 여행 계획의 모든 도시가 같은 컴포넌트에 속하는지 확인
    root = find(parent, plan[0])
    for city in plan:
        if find(parent, city) != root:
            print("NO")
            return

    print("YES")


main()

---

※ 알아야 할 것

https://gmlwjd9405.github.io/2018/08/31/algorithm-union-find.html

[알고리즘] Union-Find 알고리즘 - Heee's Development Blog