Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 간선이 최소 신장 트리(MST)에 포함되는지 확인하는 방법

2차원 리스트 edges가 하나의 무방향 그래프를 나타낸다고 가정해 봅시다. 이 리스트의 각 항목은 (u, v, w) 형태의 간선 정보를 담고 있으며, 이는 노드 u와 v가 가중치 w를 가지는 간선으로 연결되어 있음을 의미합니다. 여기에 정수 a와 b가 추가로 주어지는데, 이 값들은 간선 (a, b)를 나타냅니다. 우리가 확인해야 할 것은 바로 이 간선 (a, b)가 최소 신장 트리(Minimum Spanning Tree, MST)의 일부에 해당하는지 여부입니다.

참고: 그래프는 반드시 연결 그래프여야 하며, 간선 (a, b)는 그래프 내에 실제로 존재해야 합니다.

문제 예시

입력이 다음과 같다고 가정해 보겠습니다.

[[0, 2, 100],
[1, 2, 200],
[1, 3, 100],
[2, 3, 300]]
a = 0
b = 2

이 경우 출력 결과는 True가 됩니다.

해결 아이디어

이 문제의 핵심은 MST의 사이클 성질(Cycle Property)에 있습니다. 어떤 간선 (a, b)보다 가중치가 엄격히 작은 간선들만 사용하여 a에서 b로 도달하는 경로가 존재한다면, 해당 간선은 MST에 반드시 포함될 필요가 없습니다. 이미 더 저렴한 대체 경로가 존재하기 때문입니다. 반대로, 더 가벼운 간선들만으로는 a와 b를 연결할 수 없다면, 이 간선은 모든 최소 신장 트리에 반드시 포함되어야 합니다.

따라서 다음 단계로 문제를 해결할 수 있습니다.

1. findPath() 함수 정의

간선 목록과 두 노드 a, b를 받아 두 노드 사이에 경로가 존재하는지 재귀적으로(DFS 방식) 탐색합니다.

  • a와 b가 같으면 True를 반환합니다.
  • 간선 목록이 비어 있으면 False를 반환합니다.
  • 각 간선 x에 대해 다음을 수행합니다.
    • x의 가중치(x[2])가 -1이면 이미 사용된 간선이므로 건너뜁니다.
    • x[0]이 a와 같으면 new_a를 x[1]로 설정하고, x[1]이 a와 같으면 new_a를 x[0]으로 설정합니다.
    • new_a가 유효하면(-1이 아니면), 현재 간선 x를 목록에서 제거한 뒤 findPath(edges, new_a, b)를 재귀 호출합니다.
    • 경로가 발견되면 True를 반환하고, 그렇지 않으면 간선 x를 목록 끝에 다시 추가해 원상 복구합니다.
  • 모든 간선을 확인했는데도 경로를 찾지 못하면 False를 반환합니다.

2. 메인 로직(solve 함수)

  • 입력 배열 'edges'에서 조건 ((x[0] == a and x[1] == b) 또는 (x[1] == a and x[0] == b))를 만족하는 간선 x의 가중치를 weight 변수에 저장합니다.
  • edges를 필터링하여 가중치가 weight보다 작은 간선들만 남깁니다.
  • not findPath(edges, a, b)를 반환합니다. 즉, 더 가벼운 간선만으로 a와 b가 연결되지 않는다면 True(간선이 MST에 포함됨)를 의미합니다.

구현 코드

아래 구현을 통해 더 잘 이해해 보겠습니다.

class Solution:
    def findPath(self, edges, a, b):
       if a == b:
          return True
       if not edges:
          return False
       for x in edges:
          if x[2] == -1:
            continue
          new_a = -1
          if x[0] == a:
            new_a = x[1]
          elif x[1] == a:
            new_a = x[0]
          if new_a != -1:
            edges.remove(x)
            if self.findPath(edges, new_a, b):
              return True
            edges.append(x)
       return False

    def solve(self, edges, a, b):
       weight = next(x for x in edges if (x[0] == a and x[1] == b) or (x[1] == a and x[0] == b))[2]
       edges = [x for x in edges if x[2] < weight]
       return not self.findPath(edges, a, b)

ob = Solution()
print(ob.solve([
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
], 0, 2))

입력

[
[0, 2, 100],
[1, 2, 200],
[1, 3, 100],
[2, 3, 300]
], 0, 2

출력

True

동작 원리 정리

위 예제에서 간선 (0, 2)의 가중치는 100입니다. 가중치가 100 미만인 간선은 하나도 없으므로, 필터링 후 남은 간선 목록은 비어 있게 됩니다. 따라서 findPath는 False를 반환하고, 최종 결과는 not False 즉 True가 됩니다. 이는 간선 (0, 2)가 최소 신장 트리에 반드시 포함되어야 한다는 의미입니다.

이 알고리즘은 간선 수를 E, 노드 수를 V라고 할 때 DFS 경로 탐색이 O(E × (V + E))의 시간 복잡도를 가지므로, 그래프 크기가 작은 경우에 적합합니다. 큰 그래프에서는 유니온-파인드(Union-Find) 자료구조를 활용하면 훨씬 효율적으로 처리할 수 있습니다.