2233번: 사과나무
첫째 줄에 정점의 개수 N(1 ≤ N ≤ 2,000)이 주어진다. 둘째 줄에는 벌레가 만드는 2×N자리의 이진수가 주어진다. 셋째 줄에는 썩은 사과의 위치를 나타내는 두 정수 X, Y가 주어진다. 이는 2×N자리
www.acmicpc.net
코드
import sys
sys.setrecursionlimit(10**6)
n = int(input())
binary = '9' + input()
x, y = map(int, input().split())
parent = [[] for _ in range(n+1)] # 각 노드의 부모(가까운 노드부터)
depth = [-1] * (n+1) # 각 노드의 깊이
removed = []
inout = [[] for _ in range(n+1)]
def dfs(idx, stack, num):
if idx >= len(binary):
return
if binary[idx] == '0':
if stack:
parent[num] = [stack[-1][1]]+parent[stack[-1][1]]
depth[num] = len(stack)
stack.append((idx, num))
inout[num].append(idx)
num += 1
else:
i, n = stack.pop()
if idx == x or idx == y or i == x or i == y:
removed.append(n)
inout[n].append(idx)
dfs(idx+1, stack, num)
dfs(1, [], 1)
if len(removed) == 1:
print(*inout[removed[0]])
else:
# LCA
a, b = removed
if depth[a] > depth[b]:
a, b = b, a
while depth[a] < depth[b]:
b = parent[b][0]
if a == b:
print(*inout[a])
else:
for i in range(n):
if parent[a][i] == parent[b][i]:
print(*inout[parent[a][i]])
break
이 문제를 풀려면 우선 최소 공통 조상(LCA, Lowest Common Ancestor)에 대한 이해가 필요하다. 나는 아래의 블로그에서 힌트를 얻어서 풀 수 있었다.
42. 최소 공통 조상(Lowest Common Ancestor)
이번 시간에 공부할 내용은 최소 공통 조상(Lowest Common Ancestor)입니다. 최소 공통 조상이란...
blog.naver.com
최소 공통 조상이란 트리(Tree) 구조에서 두 노드의 공통된 조상 중에서 가장 가까운 조상을 말한다.
최소 공통 조상을 구하기 위해서는 각 노드들의 깊이와 조상들을 구하면 된다.
이를 구하기 위해서 DFS로 두번째 줄 입력으로 주어졌던 2×N 이진수를 탐색하여 각 숫자들의 노드와 그 노드의 깊이, 조상들을 구한다.
- 2×N 이진수의 idx번째 숫자가 0이라면 아래의 과정들을 수행한다.
- stack이 비어있지 않다면 이번 노드 num의 조상을 저장하기 위해 가장 가까운 조상인 stack의 top에 있는 노드와 그 노드의 조상들을 리스트로 만들어서 num의 조상 노드들로 저장.
- num의 depth를 현재 stack의 길이로 저장.
- 나중에 노드 번호가 아닌 방문할 때 순서번호와 나갈 때 순서번호를 출력해줘야 하므로 inout배열에 현재 idx를 저장.
- 노드 번호로 num이 이번에 쓰였으므로 +1.
- 2×N 이진수의 idx번째 숫자가 1인 경우
- 현재 stack의 top을 pop해서 i, n에 저장. i는 2×N 이진수의 인덱스, n은 노드 번호.
- 만약 pop된 인덱스나 현재 인덱스 idx가 x 또는 y라면 제거할 노드 번호를 저장할 removed에 n을 저장해줌.
- 나갈 때 순서번호를 기록하기 위해 inout[n]에 idx 추가.
이렇게 하여 dfs 함수를 실행하면 각 노드의 깊이, 조상들 그리고 제거할 사과의 노드 번호를 알 수 있다. 이제 이 값을 가지고 제거할 사과들의 최소 공통 조상을 구하면 된다.
만약 제거할 사과가 하나뿐이라면 조상을 찾을 필요없이 그 사과를 제거하면 된다.
아니라면 아래의 과정으로 최소 공통 조상을 찾아주면 된다.
- removed에 저장된 노드 번호 2개를 a, b에 저장하고, 둘 중 더 깊이가 깊은 것이 b가 되도록 한다.
- b의 깊이가 a와 같아지도록 b의 조상을 거슬러 올라온다.
- 높이를 맞춘 뒤, a와 b의 조상이 같은 노드일때까지 탐색하다가 같은 노드를 찾으면 그 노드의 inout 값을 출력하고 반복문을 벗어난다.
'문제풀이 > DFS_BFS' 카테고리의 다른 글
[Python/파이썬] 백준 14716번 현수막 (0) | 2024.03.20 |
---|---|
[Python/파이썬] 백준 18513번 샘터 (0) | 2024.03.13 |
[Python/파이썬] 백준 17073번 나무 위의 빗물 (0) | 2024.02.19 |
[Python/파이썬] 백준 10819번 차이를 최대로 (0) | 2024.02.08 |
[Python/파이썬] 백준 7562번 나이트의 이동 (0) | 2024.02.02 |