(3830) 교수님은 기다리지 않는다

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

단계별로 풀어보기

43단계(유니온 파인드 2) 6번째

 

 

 

유니온 파인드로 부모 정점을 의미하는 parent와 그 부모 정점과 무게 차이를 의미하는 diff, 2가지 변수를 관리해야 한다.

find를 통해 부모로 계속 올라가는 과정에서 diff가 갱신되고 루트 정점(샘플)과의 무게 차이를 구할 수 있으며, 상대 무게 차이를 이용해 같은 그룹에 속한 두 샘플의 무게 차이도 구할 수 있다.

설명을 보면 알겠지만 경로 압축을 하면 안 된다. 따라서 랭크 기반 최적화(안 하면 시간초과로 기억함)를 해야 하며 집합의 높이를 의미하는 height 변수까지 3가지 변수를 관리해야 한다.

병합 과정에서는 합쳐지는 쪽 루트 정점의 parent와 diff를 갱신해야 한다.

 

 

 

#include <iostream>
#include <vector>
using namespace std;

struct Info{
  int parent;
  int height;
  int diff;
};

void solve(int n, int m){
  vector<Info> uf(n+1);
  
  for(int i=1;i<=n;i++)  uf[i].parent= i;
  
  auto find= [&](auto self, int x, int w) -> pair<int, int> {
    if(uf[x].parent==x)  return {x, w};
    else  return self(self, uf[x].parent, w+uf[x].diff);
  };
  
  auto merge= [&](int x, int y, int w){
    auto [xr, xw]= find(find, x, 0);
    auto [yr, yw]= find(find, y, 0);
    
    if(xr==yr)  return;
    
    if(uf[xr].height<uf[yr].height){
      uf[xr].parent= yr;
      uf[xr].diff= -w-xw+yw;
    }else{
      uf[yr].parent= xr;
      uf[yr].diff= w+xw-yw;
      
      if(uf[xr].height==uf[yr].height)  uf[xr].height++;
    }
  };
  
  while(m--){
    char qtype;
    
    cin >> qtype;
    
    if(qtype=='!'){
      int a, b, w;
      
      cin >> a >> b >> w;
      
      merge(a, b, w);
    }else{
      int a, b;
      
      cin >> a >> b;
      
      auto ares= find(find, a, 0);
      auto bres= find(find, b, 0);
      
      if(ares.first==bres.first)  cout << bres.second-ares.second << "\n";
      else  cout << "UNKNOWN\n";
    }
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  
  while(true){
    int n, m;
    
    cin >> n >> m;
    
    if(!n)  break;
    
    solve(n, m);
  }
  
  return 0;
}