2. 교수님은 기다리지 않는다

문제 링크

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

3830번: 교수님은 기다리지 않는다

 

문제 선정 과정

그래프 기반 문제로, 온라인 쿼리로 간선이 추가되며 두 정점 사이의 거리를 구해야 하는 문제로 보였습니다.

난이도도 적당해 보이기 때문에 이 문제를 풀기로 하였습니다.

 

문제 풀이 전략

문제를 모델링하면 다음과 같습니다.

• $N$개의 정점에 대하여 두 가지의 쿼리가 주어집니다.
• 쿼리 1: 가중치가 있는 방향 간선이 추가됩니다.
• 쿼리 2: 두 정점이 주어질 때, 연결되어 있다면 한 정점에서 다른 정점까지의 경로를 따라갔을 때의 가중치의 합을 구합니다.

그래프 상에서 정점의 연결 관계를 파악해야 한다는 점에서 Union-Find 자료구조를 써야겠다는 생각이 들었습니다.

 

자료 구조

Union-Find 자료구조는 상호 배타적 집합(Disjoint set)을 표현할 때 쓰는 자료구조입니다.

이름에서 알 수 있듯, 임의의 서로 다른 두 집합에 대해 교집합이 공집합이 됩니다.

 

Union-Find 자료구조는 트리를 사용하여 구현합니다.

두 정점에 대해 각각의 루트 노드가 같다면 같은 집합에 있고, 다르다면 다른 집합에 있다고 판단하는 방식입니다.

 

두 집합을 합치려면, 간단히 두 집합의 루트 노드를 이어주는 방식으로 해결할 수 있습니다.

 

 

이제 이를 활용하여 문제를 해결해 봅시다.

Union-Find 자료구조를 잘 활용하기 위해서는 find 함수가 호출될 때마다 해당 정점을 바로 루트 노드로 연결시켜 줍니다.

 

이때 간선에 가중치가 주어지기 때문에, 이를 고려하여 합해주어야 합니다.

이를 구현한 코드입니다.

int find(int i)
{
  if(parent[i] == i) return i;

  int root = find(parent[i]);
  dist[i] += dist[parent[i]];
  return parent[i] = root;
}

 

다음으로, 두 노드를 연결하는 함수를 구현합니다.

방향성 때문에 헷갈리지 않도록, dist의 정의를 루트 노드보다 dist만큼 무겁다는 뜻으로 설정하겠습니다.

 

두 집합의 루트 노드를 서로 연결해 주어야 하기 때문에, 원래의 간선의 가중치를 다음과 같이 수정해야 합니다.

W = dist[y] - dist[x] + w

void merge(int x, int y, LL w)
{
  int xroot = find(x);
  int yroot = find(y);

  dist[xroot] = dist[y] - dist[x] + w;
  parent[xroot] = yroot;
}

 

마지막으로, 두 무게의 차는 루트 노드와의 편차의 차이로 구할 수 있습니다.

지금까지의 내용을 코드로 구현하면 다음과 같습니다.

#include <iostream>
#include <cstring>

using namespace std;
typedef long long int LL;

int parent[100001];
LL dist[100001];

int find(int i)
{
  if(parent[i] == i) return i;

  int root = find(parent[i]);
  dist[i] += dist[parent[i]];
  return parent[i] = root;
}

void merge(int x, int y, LL w)
{
  int xroot = find(x);
  int yroot = find(y);

  dist[xroot] = dist[y] - dist[x] + w;
  parent[xroot] = yroot;
}

int main()
{
  ios::sync_with_stdio(false);
  cin.tie(0);
  
  while(1)
    {
      int n,m;
      cin >> n >> m;
      if(n + m == 0) return 0;

      memset(dist,0,sizeof(dist));
      for(int i=1;i<=n;i++) parent[i] = i;

      while(m--)
        {
          char cmd;
          cin >> cmd;

          if(cmd == '!')
          {
            int x,y;
            LL w;
            cin >> x >> y >> w;

            if(find(x) != find(y)) merge(y,x,w);
          }
          else
          {
            int a,b;
            cin >> a >> b;

            if(find(a) != find(b)) cout << "UNKNOWN\n";
            else cout << dist[b] - dist[a] << '\n';
          }
        }
    }
}

 

느낀 점

Union-Find 알고리즘에 가중치를 추가해야 한다는 게 꽤나 복잡했던 문제입니다.

그래도 그림을 그려가보며 한번 제대로 구현하니 심각한 오류 없이 풀 수 있었습니다.