어려움MST힙그래프

도시 연결하기 (프림)

정점 개수 n과 무방향 가중 간선 목록 edges가 주어져요. 각 원소 [a, b, cost]는 a와 b를 비용 cost로 잇는 간선이에요. 모든 정점을 잇는 데 드는 최소 비용을 반환해요. 다 이을 수 없으면 -1을 반환해요.

크루스칼이 간선을 골랐다면, 프림은 정점을 하나씩 트리에 붙여 키워요. 트리에서 밖으로 나가는 간선 중 가장 싼 것을 골라 새 정점을 들여요. 다익스트라와 뼈대가 같은데, 힙에 넣는 값이 '누적 거리'가 아니라 '그 간선 하나의 비용'이라는 점만 달라요.

간선이 빽빽한 그래프에서 특히 잘 어울려요.

예시

예시 1

입력: n = 4, edges = [[0, 1, 1], [1, 2, 2], [2, 3, 3], [0, 3, 4], [0, 2, 5]]
출력: 6
설명: 0에서 시작해 1(1), 2(2), 3(3)을 차례로 붙이면 합이 6이에요.

예시 2

입력: n = 3, edges = [[0, 1, 1]]
출력: -1
설명: 2번 정점을 트리에 붙일 간선이 없어요.

제한 사항

1 ≤ n ≤ 100,000
0 ≤ edges.length ≤ 300,000
0 ≤ a, b < n, 1 ≤ cost ≤ 1,000,000

인접 리스트를 [비용, 이웃]으로 만들어요. [0, 0]을 힙에 넣고 시작해, 가장 싼 간선을 꺼내 아직 트리에 없는 정점이면 붙이고 비용을 더해요. 그 정점의 간선들을 힙에 넣어요. 붙인 정점이 n개면 합을, 힙이 먼저 비면 -1을 반환해요.

javascript
Copy code
class MinHeap {
  /** @param {(a: any, b: any) => number} [compare] 음수면 a가 먼저 나옵니다. */
  constructor(compare) {
    this.items = [];
    this.compare = compare ?? ((a, b) => (a < b ? -1 : a > b ? 1 : 0));
  }
  get size() {
    return this.items.length;
  }
  peek() {
    return this.items.length > 0 ? this.items[0] : null;
  }
  push(value) {
    const a = this.items;
    a.push(value);
    let i = a.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.compare(a[parent], a[i]) <= 0) break;
      [a[parent], a[i]] = [a[i], a[parent]];
      i = parent;
    }
  }
  pop() {
    const a = this.items;
    if (a.length === 0) return null;
    const top = a[0];
    const last = a.pop();
    if (a.length > 0) {
      a[0] = last;
      let i = 0;
      for (;;) {
        const l = i * 2 + 1;
        const r = l + 1;
        let small = i;
        if (l < a.length && this.compare(a[l], a[small]) < 0) small = l;
        if (r < a.length && this.compare(a[r], a[small]) < 0) small = r;
        if (small === i) break;
        [a[small], a[i]] = [a[i], a[small]];
        i = small;
      }
    }
    return top;
  }
}
function primMst(n, edges) {
  const graph = Array.from({ length: n }, () => []);
  for (const [a, b, cost] of edges) {
    graph[a].push([cost, b]);
    graph[b].push([cost, a]);
  }
  const inTree = new Array(n).fill(false);
  const heap = new MinHeap((a, b) => a[0] - b[0]);
  heap.push([0, 0]);
  let total = 0;
  let used = 0;
  while (heap.size > 0 && used < n) {
    const [cost, v] = heap.pop();
    if (inTree[v]) continue;
    inTree[v] = true;
    total += cost;
    used++;
    for (const [w, next] of graph[v]) {
      if (!inTree[next]) heap.push([w, next]);
    }
  }
  return used === n ? total : -1;
}

이전 문제크루스칼 MST

다음 문제서브트리의 크기

예시 테스트만 실행 (Cmd/Ctrl+Enter)

숨김 테스트까지 채점 (Cmd/Ctrl+Shift+Enter)

에디터를 불러오고 있어요…

코드를 작성하고를 눌러 예시 테스트를 확인해 보세요.