JS Coding Blog

전체 지연 합이 가장 작은 중앙 허브 찾기

자바스크립트 코딩테스트 문제로 tree-rerooting 주제를 연습해보세요. 난이도는 hard이며, 브라우저에서 바로 JavaScript로 풀이를 실행할 수 있습니다.

today hard tree-rerooting 함수명: solution 제한 시간: 500ms

가중치가 있는 트리 네트워크에서 모든 정점까지의 거리 합이 가장 작은 중앙 허브를 찾는 문제입니다.

문제 설명

정점 개수 n과 가중치가 있는 무방향 간선 목록 edges가 주어집니다. 정점 번호는 0부터 n - 1까지입니다.

각 정점을 허브로 선택했을 때, 그 정점에서 모든 다른 정점까지의 최단 거리 합을 계산할 수 있습니다. 이 거리 합이 가장 작은 정점 번호를 반환하는 solution 함수를 작성하세요.

만약 최소 거리 합을 만드는 정점이 여러 개라면, 정점 번호가 가장 작은 것을 반환합니다.

트리는 사이클이 없고 모든 정점이 연결되어 있으므로, 두 정점 사이의 경로는 항상 하나뿐입니다.

제한사항

  • 1 <= n <= 200,000
  • edges.length === n - 1
  • 각 간선은 [u, v, w] 형태입니다.
  • 0 <= u, v < n
  • u !== v
  • 1 <= w <= 1,000,000
  • 입력 그래프는 가중치가 있는 무방향 트리입니다.
  • 반환값은 조건을 만족하는 허브 정점의 번호입니다.

예시

  • 입력: n = 4, edges = [[0, 1, 1], [1, 2, 2], [1, 3, 3]] → 출력: 1
  • 입력: n = 5, edges = [[0, 1, 1], [1, 2, 1], [2, 3, 1], [3, 4, 1]] → 출력: 2
  • 입력: n = 2, edges = [[0, 1, 7]] → 출력: 0

첫 번째 예시에서 각 정점을 허브로 골랐을 때 거리 합은 다음과 같습니다.

  • 정점 0: 1 + 3 + 4 = 8
  • 정점 1: 1 + 2 + 3 = 6
  • 정점 2: 2 + 3 + 5 = 10
  • 정점 3: 3 + 4 + 5 = 12

따라서 정답은 1입니다.

힌트

  • 모든 정점에서 DFS/BFS를 다시 돌리면 너무 느립니다.
  • 먼저 아무 정점 하나를 루트로 잡고, 그 루트에서 모든 정점까지의 거리 합을 구해 보세요.
  • 루트를 부모에서 자식으로 옮길 때, 가까워지는 정점 수와 멀어지는 정점 수를 나눠 보면 거리 합 갱신 규칙이 보입니다.
  • 각 자식의 서브트리 크기를 알고 있으면 다음 루트의 거리 합을 상수 시간에 계산할 수 있습니다.

해설

핵심은 모든 정점을 각각 시작점으로 다시 계산하지 않고, 한 정점에서 구한 정보를 다른 정점으로 전달하는 것입니다.

트리는 경로가 하나뿐이라서, 루트를 한 칸 옮길 때 거리 합이 어떻게 바뀌는지 정확히 계산할 수 있습니다.

1) 임의의 루트에서 시작하기

우선 0번 정점을 루트라고 생각합시다.

첫 번째 DFS(또는 스택 기반 순회)에서 아래 두 가지를 구합니다.

  • subtreeSize[x]: 정점 x를 루트로 하는 서브트리의 정점 수
  • distSum[0]: 루트 0에서 모든 정점까지의 거리 합

distSum[0]은 루트에서 각 정점까지의 깊이(가중 거리)를 더하면 됩니다.

2) 루트를 이웃으로 옮길 때의 변화

이제 부모 u에서 자식 v로 루트를 옮긴다고 생각해 봅시다. 간선 가중치를 w라고 하면:

  • v의 서브트리에 있는 subtreeSize[v]개 정점은 새 루트 v에 w만큼 더 가까워집니다.
  • 나머지 n - subtreeSize[v]개 정점은 새 루트 v에 w만큼 더 멀어집니다.

그래서 거리 합은 이렇게 바뀝니다.

 distSum[v] = distSum[u] - subtreeSize[v] * w + (n - subtreeSize[v]) * w

정리하면:

 distSum[v] = distSum[u] + (n - 2 * subtreeSize[v]) * w

이 식 덕분에, 부모의 정답만 알면 자식의 정답을 O(1)에 구할 수 있습니다.

3) 두 번째 순회로 모든 정점 답 구하기

두 번째 DFS에서는 위 공식을 이용해 distSum[x]를 모든 정점에 대해 채웁니다.

그러면 각 정점을 허브로 삼았을 때의 전체 거리 합을 전부 구할 수 있고, 그중 가장 작은 값을 갖는 정점을 고르면 됩니다.

동률이면 번호가 가장 작은 정점을 반환해야 하므로, 비교할 때:

  1. 거리 합이 더 작으면 갱신
  2. 거리 합이 같으면 정점 번호가 더 작은 쪽으로 갱신

하면 됩니다.

4) 왜 빠른가?

  • 인접 리스트 구성: O(n)
  • 첫 번째 순회: O(n)
  • 두 번째 순회: O(n)

총 O(n)에 해결할 수 있습니다.

참고 구현 아이디어

재귀 DFS는 입력이 매우 클 때 호출 스택이 부족할 수 있으므로, JavaScript에서는 반복문 + 스택 방식이 더 안전합니다.

구현 흐름은 보통 이렇게 잡을 수 있습니다.

  1. 인접 리스트 구성
  2. 스택으로 부모/방문 순서 기록
  3. 방문 순서를 뒤에서부터 보며 subtreeSize 계산
  4. 루트 0의 거리 합 계산
  5. 다시 순회하며 rerooting 공식으로 모든 distSum 계산
  6. 최소 거리 합을 만드는 정점 선택

이 문제는 트리 DP에서 자주 나오는 rerooting 감각을 익히기 좋은 문제입니다. 한 루트에서 구한 값을 다른 루트의 답으로 밀어 주는 방식이 핵심입니다.

코드 작성

starter code를 바탕으로 함수를 완성한 뒤 예제 테스트를 실행해보세요.

JavaScript 에디터 로딩 중...

커스텀 테스트

함수 인자를 JSON 배열 형태로 입력하세요. 예: [3, 5], [[1, 2, 3]]

아직 실행하지 않았습니다.

실행 결과

아직 실행하지 않았습니다.

예제 테스트를 실행하면 여기에서 결과를 확인할 수 있습니다.

댓글

문제 풀이 아이디어, 질문, 반례를 자유롭게 나눠보세요.