[코테 스터디] 그래프 이론, 여행 계획

Q41. 여행 계획

🐣문제

한울이가 사는 나라에는 N개의 여행지가 있으며, 각 여행지는 1~N번까지의 번호로 구분됩니다. 또한 임의의 두 여행지 사이에는 두 여행지를 연결하는 도로가 존재할 수 있습니다. 이때, 여행지가 도로로 연결되어 있다면 양방향으로 이동이 가능하다는 의미입니다. 한울이는 하나의 여행 계획을 세운 뒤에 이 여행 계획이 가능한지의 여부를 판단하고자 합니다.


입력 조건
첫째 줄에 여행지의 수 N과 여행 계획에 속한 도시의 수 M이 주어집니다.
다음 N개의 줄에 걸쳐 NxN 행렬을 통해 임의의 두 여행지가 서로 연결되어 있는지의 여부가 주어집니다. 그 값이 1이라면 서로 연결되었다는 의미이며, 0이라면 서로 연결되어 있지 않다는 의미입니다.
마지막 줄에 한울이의 여행 계획에 포함된 여행지의 번호들이 주어지며 공백으로 구분합니다.

입력 예시
5 4
0 1 0 1 1
1 0 1 1 0
0 1 0 0 0
1 1 0 0 0
1 0 0 0 0
2 3 5 4

5 4
0 1 0 1 0
1 0 1 1 0
0 1 0 0 0
1 1 0 0 0
0 0 0 0 0
2 3 5 4


출력 조건
첫째 줄에 한울이의 여행 계획이 가능하다면 YES를, 불가능하다면 NO를 출력합니다.

출력 예시
YES

NO

🐥풀이

간단히 요약하면, 이 문제는 한울이의 여행 계획에 속한 도시가 모두 연결되어 있는 지를 알아보아야 하는 문제이다.

쉽게 말해, 그래프 이론을 활용하여 노드의 루트를 묶어 서로 연결된 노드끼리 같은 집합으로 묶고, 여행 계획에서 노드 연결 유무를 확인하는 방식으로 풀면 된다.


팀 합치기
먼저, a 노드와 b 노드의 팀을 합치는 코드부터 알아보자. 서로 연결된 노드끼리는 같은 집합으로 묶어주는 과정이다.

def union_team(team, a, b):
  a = check_team(team, a) # a노드의 루트 번호 확인
  b = check_team(team, b) # b노드의 루트 번호 확인

  if a<b:
    team[b]=a
  else:
    team[a]=b

각 노드의 루트 노드를 체크하고, 루트가 다르면 루트 번호가 작은 번호로 루트를 통합시킨다.


루트 확인하기
다음으로, 특정 노드의 루트를 확인하는 코드를 알아보자. 팀 합치기에서 루트를 구하기 위해 필요한 함수이다.

# 루트 확인
def check_team(team, x):
  if team[x]!=x:
    team[x] = check_team(team, team[x])
  return team[x]

재귀를 통해 루트를 확인할 수 있다. (팀 합치기 전, 초기 리스트는 team[i]=i로 미리 초기화 되어 있음) 타고 타고 올라가서 루트의 번호가 노드 자신의 번호와 일치하면 루트 노드의 특징을 지니고 있는 것이다.


위 두 함수를 사용하여, 해당 문제를 빠르게 풀이할 수 있다.
먼저 입력값을 받아와 1이면 연결이라고 했으므로, 1이면 두 도시의 집합을 같은 집합으로 묶는다. (union_team)
그 후, 마지막 줄의 한울이의 여행 계획을 리스트로 들고 와서 속한 모든 도시의 집합(루트)이 같은 집합인 지를 확인하고, 어느 하나라도 다른 집합이면 "NO"를, 모두 같은 집합이면 "YES"를 출력한다.

🐓코드

# 여행지수, 도시수
n, m = map(int, input().split())  

# 팀 초기화
team = [0]*(n+1)
for i in range(1,n+1): team[i] = i

# 루트 확인
def check_team(team, x):
  if team[x]!=x:
    team[x] = check_team(team, team[x])
  return team[x]

# 팀 합치기
def union_team(team, a, b):
  a = check_team(team, a)
  b = check_team(team, b)
  if a<b:
    team[b]=a
  else:
    team[a]=b

# 여행지 연결 가능하면 합치기
for i in range(n):
  dest = list(map(int, input().split()))
  for j in range(n):
    if dest[j]==1:
      union_team(team, i+1, j+1)

# 한울이의 여행 계획
plan = list(map(int, input().split()))


# center은 첫여행지 루트, 첫여행지 root랑 다른 여행지 있으면 도달X 없으면 도달 O
def possible(center):
  for p in plan:
    if center!=check_team(team, p):
      return "NO"
  return "YES"

# 여행 가능 여부 출력
print(possible(check_team(team, plan[0])))

⭐2022.05.11

스터디 끝무리의 단원 회고 시작..!! 팀 합치기랑 루트 확인하기 함수 기억이 몽실몽실 떠오르는 중 ~.~