[LeetCode] 3600. Maximize Spanning Tree Stability with Upgrades (Hard)

이 문제를 풀기 위해서는 MST(Minimum Spanning Tree, 최소 신장 트리)에 대해서 알아야 한다.

Spanning Tree(신장 트리)는 그래프 내의 모든 정점을 포함하는 트리이다.

스패닝 트리는 다음과 같이 여러 개가 나올 수 있다.

스패닝 트리는 사이클을 만들면 안 되고 n개의 노드를 (n-1)개의 간선으로 연결한다.

그리고 이 스패닝 트리를 연결하는 간선에 가중치(strength)가 있을 때 그 가중치의 값을 최소로 구하는 경우가 바로 MST이다.

1. 접근 : 문제를 단순화 하기

이 문제에서는 노드가 n개 주어지고 각각의 간선이 edges[i]로 주어진다고 했다. 

edges[i] = [u, v, s, must] 인데 u - v 이어지는 무방향 간선이며 여기에 가중치가 s, 그리고 must가 1이면 필수, 0이면 한 번까지 업그레이드가 가능하다. 업그레이드를 하면 가중치는 2배가 된다.

이 노드와 간선들의 그래프에서 MST를 구하고 그 MST들의 최대 가중치 값을 구하는 것이 문제이다.

n <= 10^5, u,v < 10^5, s <= 10^5 이다.

즉 이 문제는 시간 복잡도를 O(N^2) 이상으로 푸는 순간 바로 시간 초과다. 이 점을 기억하고 문제를 접근해야 한다.

2. must 간선을 먼저 공략하자

edges[i]에서 must = 1 이면 그 가중치는 변하지 않고 반드시 들어가야 한다.

따라서 must = 1 인 간선들을 먼저 찾아서 연결해 본다. 만약 이 때 사이클이 생기면 -1을 리턴해야 한다.

2-1. DSU(Disjoint Set Union, 서로소 집합 자료구조)

이 문제를 풀기 위해서는 DSU의 개념에 대해서 먼저 알아야 한다.

서로소 집합은 두 집합이 겹치는 원소가 없는 집합 관계를 의미한다. 예를 들어 {1,3,5}와 {2,4}는 서로소 집합 관계이다. 

우리가 어떤 원소가 해당 집합에 속해있는지를 알기 위해서 이 DSU 자료구조를 사용한다.

DSU에는 find와 union 이라는 두 가지의 중요한 연산이 있다.

• find(x) : 원소 x가 어느 집합에 속해 있는지 대표(root)를 찾아 반환
• union(a, b) : 원소 a가 속한 집합과 원소 b가 속한 집합을 결합

이 개념이 다소 헷갈릴 수가 있어서 보충 설명을 하면 각각의 서로소 집합에는 하나씩 대표 노드가 존재한다. 그래서 그 집합에 포함된 요소들 중 하나를 인풋으로 find() 메서드에 넣으면 대표 노드를 아웃풋으로 반환하는 것이다.

그리고 어떤 두 요소가 각각 속해있는 집합이 있다고 했을 때, 그 두 요소를 union() 메서드에 인풋으로 받으면, 그 두 집합의 합집합을 아웃풋으로 반환한다. 아래의 그림을 참고하면 된다.

find 메서드와 union 메서드

우리는 이러한 그래프를 트리 형태로 작성할 수 있는데, 하나의 트리에서 어떤 요소가 대표 노드를 찾아가는 과정을 담은 배열이 parent 이다. parent[i] = j 라고 하는 것은 i 요소의 부모 요소가 j라는 말이고 이렇게 타고타고 올라가면 대표 노드를 찾을 수 있다.

이러한 DSU 자료구조를 타입스크립트로 작성하면 다음과 같이 작성할 수 있다.

class DSU {
  parent: number[];
  size: number[];
  components: number;

  constructor(n: number) {
    this.parent = Array.from({ length: n }, (_, i) => i);
    this.size = Array(n).fill(1);
    this.components = n;
  }

  find(x: number): number {
    if (this.parent[x] !== x) {
      this.parent[x] = this.find(this.parent[x]);
    }
    return this.parent[x];
  }

  union(a: number, b: number): boolean {
    let pa = this.find(a);
    let pb = this.find(b);

    if (pa === pb) return false;

    if (this.size[pa] < this.size[pb]) {
      [pa, pb] = [pb, pa];
    }

    this.parent[pb] = pa;
    this.size[pa] += this.size[pb];
    this.components--;
    return true;
  }
}

이렇게 DSU 자료구조를 구현한 다음에 우리가 위에서 살펴본 must 간선을 체크해 볼 수 있다.

모든 간선들을 돌면서 must 간선일 경우(must === 1) 이 간선들의 최소 가중치를 찾고, 그 과정에서 사이클이 발견이 된다면 -1을 반환하게 한다. 타입스크립트 코드는 다음과 같이 작성할 수 있다.

function maxStability(
  n: number,
  edges: number[][],
  k: number
): number {
  let hasMust = false;
  let mustMin = Infinity;

  // 1) must 간선끼리 사이클 체크
  const mustDsu = new DSU(n);
  for (const [u, v, s, must] of edges) {
    if (must === 1) {
      hasMust = true;
      mustMin = Math.min(mustMin, s);
      if (!mustDsu.union(u, v)) {
        return -1;
      }
    }
  }
}

 

3. 모든 간선을 체크해서 그래프 연결 가능 여부를 확인한다

그 다음으로는 모든 간선들을 다 연결해서 그래프를 연결할 수 있는지를 체크해야 한다. 새로운 DSU 인스턴스를 만들어서 모든 노드들을 union 메서드로 연결시킨다. 그리고 DSU의 component라는 변수는 지금 집합이 총 몇 개인지를 의미하는데, 처음 생성자 객체 DSU(n)에서는 이 값이 n이고, union이 하나씩 생길 때 마다 두 개의 집합이 하나가 되므로 -1을 해 주게 된다. 만약 모든 노드를 결합했는데 이 값이 1이 아니라면 어디선가 끊어져 있다는 뜻이고, 이는 그래프 연결이 이루어지지 않음을 의미해서 -1을 반환한다.

function maxStability(
  n: number,
  edges: number[][],
  k: number
): number {
  // 1) must 간선끼리 사이클 체크
  // ...

  // 2) 전체 그래프 연결 가능 여부 체크
  const allDsu = new DSU(n);
  for (const [u, v] of edges) {
    allDsu.union(u, v);
  }
  if (allDsu.components !== 1) {
    return -1;
  }
  
  // ...
}

 

4. 다음으로 안정성을 체크한다

불가능한 케이스를 제거했으면 이제 최대 가중치 값을 구해야 한다.

이 MST에서 어떤 간선 (u, v, s)가 안정성 X를 만족하려면 s >= X여야 한다. 만약 s가 X보다 작더라도 업그레이드가 가능하고 이 때 2*s >= X 이면 가능하다.

즉, 간선을 3종류로 나눌 수 있다.

1. 바로 사용 가능 : s >= X
2. 업그레이드 하면 가능 : s < X 이고 2*s >= X
3. 불가능 2*s < X

이제 결국 간선들을 골라서 모든 정점을 연결할 수 있는지를 빠르게 확인하면 된다. 여기서 이분 탐색과 Union-Find(DSU)를 사용한다.

5. 판정(check) 하기

만약 X = 6을 검사한다고 했을 때, 간선이 이렇게 있다고 가정하자. 업그레이드는 최대 k = 2

• A --- (8) --- B
• B --- (5) --- C
• A --- (3) --- C
• C --- (6) --- D
• B --- (4) --- D

1단계 : s >= 6인 간선 먼저 다 붙이기

• A --- (8) --- B
• C --- (6) --- D

현재 연결은 [A B] [C D] 이렇게 되어 있다.

2단계 : 업그레이드 하면 6 이상 되는 간선 보기

• 5 -> 10 : 가능
• 4 -> 8 : 가능
• 3 -> 6 : 가능

이제 이 간선들 중에서 서로 다른 컴포넌트를 이어주는 간선을 쓰면 된다.

예를 들어 B --- (5) --- C 를 쓰면

[A B] --- upgrade --- [C D] 

이렇게 전체 연결을 완료할 수 있다.

  const can = (x: number): boolean => {
    const dsu = new DSU(n);

    // must 간선은 반드시 포함 + 업그레이드 불가
    for (const [u, v, s, must] of edges) {
      if (must === 1) {
        if (s < x) return false;
        if (!dsu.union(u, v)) return false;
      }
    }

    // 업그레이드 없이 x 이상인 일반 간선 먼저 사용
    for (const [u, v, s, must] of edges) {
      if (must === 0 && s >= x) {
        dsu.union(u, v);
      }
    }

    // 업그레이드해서 x 이상이 되는 간선 사용
    let upgradesLeft = k;
    for (const [u, v, s, must] of edges) {
      if (must === 0 && s < x && s * 2 >= x && upgradesLeft > 0) {
        if (dsu.union(u, v)) {
          upgradesLeft--;
        }
      }
    }

    return dsu.components === 1;
  };

질문1 : 어떤 간선을 업그레이드 해야 하나?

check(X)를 할 때

• 이미 X 이상인 간선은 바로 쓰면 됨
• 업그레이드가 필요한 간선은 컴포넌트를 실제로 합칠 때만 카운트
• DSU로 합쳐지지 않는 간선은 무시

그렇기 때문에 업그레이드 간선을 순회하면서

• 2s >= X
• 아직 다른 컴포넌트를 연결한다면
• 업그레이드 1회 사용

이렇게 하면 된다.

질문2 : must 간선은 어떻게 판정하나?

must 간선은 무조건 포함이 되어야 하고 업그레이드도 할 수 없다.

따라서 must 간선의 가중치값의 최소를 mn 이라고 하면 정답은 절대 mn을 넘을 수 없다. 그래서 이분 탐색의 상한을 mn으로 둘 수 있다.

예를 들어, must 간선의 가중치가 3, 5, 7 이면 이분 탐색의 상한은 3인 것이다.

function maxStability(
  n: number,
  edges: number[][],
  k: number
): number {
  let hasMust = false;
  let mustMin = Infinity;

  // 1) must 간선끼리 사이클 체크
  // ...

  // 2) 전체 그래프 연결 가능 여부 체크
  // ...

  // upper bound
  let hi = 0;
  if (hasMust) {
    hi = mustMin;
  } else {
    for (const [, , s] of edges) {
      hi = Math.max(hi, s * 2);
    }
  }

  const can = (x: number): boolean => {
    // ...
  };

  let lo = 1;
  while (lo < hi) {
    const mid = Math.floor((lo + hi + 1) / 2);
    if (can(mid)) {
      lo = mid;
    } else {
      hi = mid - 1;
    }
  }

  return lo;
}

 

6. 최종 구현

결과적으로 최종 구현은 다음과 같다.

Step 1. must 간선만 먼저 DSU에 넣기

• 합치다가 이미 같은 집합이면 사이클 → -1
• 동시에 must 간선 최소 strength mn 기록

Step 2. 전체 그래프가 연결 가능한지 확인

• 모든 간선을 DSU에 넣어 봤는데도 하나로 안 합쳐지면 -1

Step 3. 이분 탐색

• low = 1
• high = mn (must 간선이 있으면)
• 중간값 mid = X에 대해 check(X) 수행

Step 4. check(X)

• DSU 새로 생성
• s >= X 인 간선들 먼저 union
• 그 다음 2s >= X 인 간선들 중 아직 다른 컴포넌트를 잇는 간선을 업그레이드해서 union
• 업그레이드 횟수는 최대 k
• 마지막에 전체가 하나로 연결되면 성공

class DSU {
  parent: number[];
  size: number[];
  components: number;

  constructor(n: number) {
    this.parent = Array.from({ length: n }, (_, i) => i);
    this.size = Array(n).fill(1);
    this.components = n;
  }

  find(x: number): number {
    if (this.parent[x] !== x) {
      this.parent[x] = this.find(this.parent[x]);
    }
    return this.parent[x];
  }

  union(a: number, b: number): boolean {
    let pa = this.find(a);
    let pb = this.find(b);

    if (pa === pb) return false;

    if (this.size[pa] < this.size[pb]) {
      [pa, pb] = [pb, pa];
    }

    this.parent[pb] = pa;
    this.size[pa] += this.size[pb];
    this.components--;
    return true;
  }
}

function maxStability(
  n: number,
  edges: number[][],
  k: number
): number {
  let hasMust = false;
  let mustMin = Infinity;

  // 1) must 간선끼리 사이클 체크
  const mustDsu = new DSU(n);
  for (const [u, v, s, must] of edges) {
    if (must === 1) {
      hasMust = true;
      mustMin = Math.min(mustMin, s);
      if (!mustDsu.union(u, v)) {
        return -1;
      }
    }
  }

  // 2) 전체 그래프 연결 가능 여부 체크
  const allDsu = new DSU(n);
  for (const [u, v] of edges) {
    allDsu.union(u, v);
  }
  if (allDsu.components !== 1) {
    return -1;
  }

  // upper bound
  let hi = 0;
  if (hasMust) {
    hi = mustMin;
  } else {
    for (const [, , s] of edges) {
      hi = Math.max(hi, s * 2);
    }
  }

  const can = (x: number): boolean => {
    const dsu = new DSU(n);

    // must 간선은 반드시 포함 + 업그레이드 불가
    for (const [u, v, s, must] of edges) {
      if (must === 1) {
        if (s < x) return false;
        if (!dsu.union(u, v)) return false;
      }
    }

    // 업그레이드 없이 x 이상인 일반 간선 먼저 사용
    for (const [u, v, s, must] of edges) {
      if (must === 0 && s >= x) {
        dsu.union(u, v);
      }
    }

    // 업그레이드해서 x 이상이 되는 간선 사용
    let upgradesLeft = k;
    for (const [u, v, s, must] of edges) {
      if (must === 0 && s < x && s * 2 >= x && upgradesLeft > 0) {
        if (dsu.union(u, v)) {
          upgradesLeft--;
        }
      }
    }

    return dsu.components === 1;
  };

  let lo = 1;
  while (lo < hi) {
    const mid = Math.floor((lo + hi + 1) / 2);
    if (can(mid)) {
      lo = mid;
    } else {
      hi = mid - 1;
    }
  }

  return lo;
}

시간복잡도

• can(x) 한 번: O(m α(n))
• 이분 탐색 포함 전체: O(m α(n) log V)

m <= 1e5에서도 충분히 통과 가능한 방식

 

참고자료

https://gmlwjd9405.github.io/2018/08/28/algorithm-mst.html

[알고리즘] 최소 신장 트리(MST, Minimum Spanning Tree)란 - Heee's Development Blog

https://devkuk.tistory.com/27

[자료구조] 서로소 집합 자료구조 (Disjoint Set Union, DSU) - Union Find

https://www.youtube.com/watch?v=ayW5B2W9hfo