문제 소개
정수 k와 n개의 노드로 이루어진 트리가 주어졌을 때, 서로 다른 두 정점 사이의 거리가 정확히 k가 되는 고유한 정점 쌍의 개수를 세는 것이 이 글의 목표입니다.
예를 들어 k = 2이고 다음과 같은 트리가 주어진다고 가정해 보겠습니다.
이 트리에서 거리가 정확히 2인 정점 쌍은 총 4개이므로, 기대되는 출력값은 4입니다.
알고리즘 접근 방법
이 문제는 깊이 우선 탐색(DFS)을 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 정점 v에 대해 "v로부터 거리가 d인 정점의 개수"를 저장하고, 자식 서브트리의 정보를 부모 정점으로 병합하는 과정에서 거리가 k가 되는 쌍을 세는 것입니다.
사용되는 주요 변수와 함수는 다음과 같습니다.
- N = 5005: 처리 가능한 최대 노드 수
- graph: 크기 N의 인접 리스트
- vertex_count: 각 정점에서 거리별로 도달 가능한 정점 수를 저장하는 2차원 배열
- res: 조건을 만족하는 정점 쌍의 개수(결과값)
insert_edge(x, y) 함수
x와 y 사이에 무방향 간선을 연결합니다. graph[x]의 끝에 y를 추가하고, graph[y]의 끝에 x를 추가합니다.
dfs(v, parent) 함수
- vertex_count[v][0]을 1로 설정합니다. 자기 자신까지의 거리는 0이기 때문입니다.
- v의 모든 인접 정점 i에 대해, i가 부모가 아니라면 다음 작업을 수행합니다.
- dfs(i, v)를 호출하여 자식 서브트리를 먼저 탐색합니다.
- j를 1부터 k+1까지 순회하며
res += vertex_count[i][j-1] * vertex_count[v][k-j]를 계산합니다. 자식 서브트리 안의 정점과 기존에 누적된 정점을 잇는 경로 길이가 정확히 k가 되는 쌍을 세는 단계입니다. - j를 1부터 k+1까지 순회하며
vertex_count[v][j] += vertex_count[i][j-1]로 자식 서브트리의 거리별 정보를 현재 정점에 병합합니다.
구현 예제
다음은 위 알고리즘을 파이썬으로 구현한 전체 코드입니다.
N = 5005
graph = [[] for i in range(N)]
vertex_count = [[0 for i in range(505)] for i in range(N)]
res = 0
def insert_edge(x, y):
graph[x].append(y)
graph[y].append(x)
def dfs(v, parent):
global res
vertex_count[v][0] = 1
for i in graph[v]:
if (i != parent):
dfs(i, v)
for j in range(1, k + 1):
res += vertex_count[i][j - 1] * vertex_count[v][k - j]
for j in range(1, k + 1):
vertex_count[v][j] += vertex_count[i][j - 1]
k = 2
insert_edge(1, 2)
insert_edge(2, 3)
insert_edge(3, 4)
insert_edge(2, 5)
dfs(1, 0)
print(res)입력
k = 2 insert_edge(1, 2) insert_edge(2, 3) insert_edge(3, 4) insert_edge(2, 5)
출력
4
마무리
이 알고리즘은 각 간선을 한 번씩 처리하면서 거리별 정점 개수를 병합하므로, 전체 시간 복잡도는 O(n × k)입니다. 트리 문제에서 자식 서브트리의 정보를 부모 정점으로 병합하는 이러한 패턴은 다양한 경로 계산 문제에 널리 응용되므로, 잘 익혀두면 코딩 테스트나 알고리즘 문제 해결에 큰 도움이 됩니다.