거리 제한 안의 정류장 쌍 개수

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

today hard centroid-decomposition 함수명: countStationPairsWithinDistance 제한 시간: 1200ms

가중 트리에서 서로 이동 거리가 제한 이하인 정류장 쌍의 개수를 구하세요.

문제 설명

정류장 n개와 정류장을 잇는 양방향 도로 목록 edges, 거리 제한 maxDistance가 주어집니다.

각 도로는 [a, b, distance] 형태이며, 모든 정류장은 사이클 없이 연결되어 있습니다. 두 정류장 사이의 이동 거리는 두 정류장을 잇는 유일한 경로에 포함된 도로 거리의 합입니다.

서로 다른 두 정류장의 이동 거리가 maxDistance 이하인 정류장 쌍의 총개수를 반환하는 countStationPairsWithinDistance 함수를 작성하세요. 정류장 쌍의 순서는 구분하지 않습니다.

제한사항

1 <= n <= 100,000
edges.length = n - 1
각 원소는 [a, b, distance] 형태입니다.
1 <= a, b <= n, a != b
1 <= distance <= 1,000,000
입력 그래프는 모든 정류장이 연결된 트리입니다.
0 <= maxDistance <= 100,000,000,000
반환값은 JavaScript의 안전한 정수 범위 안에 들어온다고 가정합니다.

예시

입력: 5, [[1, 2, 2], [1, 3, 3], [2, 4, 2], [2, 5, 4]], 5 -> 출력: 6
입력: 4, [[1, 2, 1], [2, 3, 1], [3, 4, 1]], 2 -> 출력: 5
입력: 5, [[1, 2, 3], [1, 3, 3], [1, 4, 3], [1, 5, 3]], 6 -> 출력: 10
입력: 1, [], 100 -> 출력: 0

힌트

한 정류장을 기준으로 모든 거리 쌍을 세면 O(n^2)이 될 수 있습니다.
트리를 균형 있게 나누는 센트로이드를 선택하면 각 정류장은 전체 분할 과정에서 O(log n)번만 다뤄집니다.
센트로이드를 지나는 쌍만 세려면 전체 거리 목록에서 만든 쌍 수에서 같은 인접 서브트리 안에서 만든 쌍 수를 빼 보세요.

해설

센트로이드 분할을 사용하면 아직 처리하지 않은 트리 조각에서 정점 수가 절반 이하가 되도록 재귀적으로 나눌 수 있습니다.

현재 트리 조각의 센트로이드를 center라고 하겠습니다. center에서 각 정점까지의 거리를 모아 정렬하면, 합이 maxDistance 이하인 두 거리의 쌍 수를 투 포인터로 셀 수 있습니다. 거리 0인 center도 목록에 포함하므로 center와 다른 정점으로 이루어진 쌍도 함께 계산됩니다.

하지만 센트로이드의 같은 인접 서브트리에 속한 두 정점은 서로 가는 경로가 center를 지나지 않습니다. 이런 쌍은 이후 그 서브트리를 분할할 때 다시 계산되므로, 현재 단계에서는 빼야 합니다.

각 단계의 계산은 다음과 같습니다.

현재 트리 조각의 센트로이드를 찾습니다.
center에서 현재 조각의 모든 정점까지 거리를 모아 유효한 쌍 수를 더합니다.
center의 각 인접 서브트리별로 거리 목록을 따로 모아 유효한 쌍 수를 뺍니다.
center를 제거한 각 트리 조각에 같은 과정을 반복합니다.

정렬된 거리 배열에서 합이 제한 이하인 쌍 수는 양끝 투 포인터로 셉니다. 왼쪽과 오른쪽 거리의 합이 제한 이하라면, 현재 왼쪽 거리와 오른쪽 이하의 모든 거리가 유효하므로 right - left를 한 번에 더하고 왼쪽을 이동합니다. 합이 너무 크면 오른쪽을 줄입니다.

각 분할 단계에서 거리 수집과 정렬을 수행하며 트리 높이는 O(log n) 수준으로 제한됩니다. 전체 시간 복잡도는 O(n log^2 n), 공간 복잡도는 O(n)입니다.

코드 작성

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

JavaScript 에디터 로딩 중...

커스텀 테스트

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

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

실행 결과