[Python] 백준 23324번: 어려운 모든 정점 쌍 최단 거리

https://www.acmicpc.net/problem/23324

23324번: 어려운 모든 정점 쌍 최단 거리

 

 > 접근

예제 입력 1에서 6이라는 결과가 어떻게 나오는지 일일히 그려본 결과, 2-3으로 분리되는 두 그룹(각각 [1,2]와 [3,4,5])에서의 노드 갯수를 곱하면 정답이 나온다는 것을 알게 되었다. 또한 1-4와 같은 간선이 추가되어 두 그룹 간의 경계가 모호해질 경우에는 결과가 0이라는 것을 알 수 있었다.

따라서 간선의 가중치가 1인 간선에 해당되는 정점 쌍(node1, node2)에서의 그룹 크기를 구하면 이 문제를 풀 수 있다.

 

 > 풀이

1. node1, node2에서 각각 DFS를 수행한다. 이 때 당연히 가중치가 1인 간선은 탐색해선 안 된다. 여기서 탐색을 1번만 해도 되는데, 만약 그 결과(count)가 N과 같다면 0을 출력하고, 아니라면 (N-count)를 곱해주면 되기 때문이다.

2. 유니온파인드를 통해 풀 수도 있다. 이 경우가 더 효율적이며 코드도 단순해진다. 입력을 받을 때, K번째 간선을 제외하고 모두 union해주면 된다.

# -*- coding: utf-8 -*-
import sys                       
input = sys.stdin.readline
    
def find(n):
  if(parent[n] == n):
    return n
  else:
    return find(parent[n])

def union(n1, n2):
  p1 = find(n1)
  p2 = find(n2)

  if(p1 < p2):
    parent[p2] = parent[p1]
  else:
    parent[p1] = parent[p2]

def get_set_count(parent_var):
  res = 0
  for i in range(1, N+1):
    if(find(i) == parent_var):
      res += 1

  return res

if __name__ == '__main__':
  N, M, K = map(int, input().split())     
  
  parent = [i for i in range(N+1)]
  node1, node2 = None, None # 가중치가 1인 정점 쌍
  for i in range(M):
    u, v = map(int, input().split())
           
    if(i != K-1):
      union(u, v)
    else:
      node1, node2 = u, v

  if(find(node1) == find(node2)):
    print(0)
  else:
    count = get_set_count(parent[1])
    print(count * (N-count))