Union-Find木(3題)

プログラミングコンテストチャレンジブック [第2版] ~問題解決のアルゴリズム活用力とコーディングテクニックを鍛える~
この本(蟻本)の章末で紹介されている練習問題をひたすら解いている。
今回は、Union-Find木を使う問題。
Union-Find木に関するいちばん面白い事実は、うまく実装されたUnion-Find木に対する操作は、計算時間がアッカーマン関数逆関数程度になるってこと。アッカーマン関数って巨大数論とかで出てくるアレだよね。信じられない速度で急激に増加する関数。だから、その逆関数ってことは計算時間めっちゃ速いってことだね。

POJ2236: Wireless Network

2236 -- Wireless Network
単純なUnion-Find木の問題。これといって工夫なし。

#include <iostream>
#include <cmath>
#include <vector>
#include <queue>
#include <utility>
#include <algorithm>
using namespace std;

//--------Union-Find木---------
vector<int> par; // 親
vector<int> rnk; // 深さ

// Union-Find木の初期化
void init(int n) {
  par  = vector<int>(n);
  rnk = vector<int>(n);
  for (int i = 0; i < n; i++) {
    par[i] = i;
    rnk[i] = 0;
  }
}

// 木の根を求める
int find(int x) {
  if (par[x] == x) {
    return x;
  } else {
    return par[x] = find(par[x]);
  }
}

// xとyの属する集合を併合
void unite(int x, int y) {
  x = find(x);
  y = find(y);
  if (x == y) return;

  if (rnk[x] < rnk[y]) {
    par[x] = y;
  } else {
    par[y] = x;
    if (rnk[x] == rnk[y]) rnk[x]++;
  }
}

// xとyが同じ集合に属するか否か
bool same(int x, int y) {
  return find(x) == find(y);
}
//-------------------------

int main() {
  // 入力
  int N, d;
  cin >> N >> d;
  d *= d; // 距離は2乗しておく
  vector<int> x(N), y(N);
  for (int i = 0; i < N; i++) {
    cin >> x[i] >> y[i];
  }
  // 修理されたコンピュータ
  vector<bool> repaired(N, false);
  // Union-Find木の初期化
  init(N);

  // 操作の読み込みはファイル終端までループ
  char op;
  while (cin >> op) {
    if (op == 'O') {
      int r;
      cin >> r;
      repaired[--r] = true;
      // union-find木の合併
      for (int i = 0; i < N; i++) {
        if (repaired[i] && (x[r] - x[i])*(x[r] - x[i]) + (y[r] - y[i])*(y[r] - y[i]) <= d) {
          unite(r, i);
        }
      }

    } else if (op == 'S') {
      int a, b;
      cin >> a >> b;
      a--; b--;
      if (same(a, b)) {
        cout << "SUCCESS" << endl;
      } else {
        cout << "FAIL" << endl;
      }
    }
  }
}

POJ1703: Find them, Catch them

1703 -- Find them, Catch them

問題概要

町に2つのギャング組織がある。N人の犯罪者について、M個の情報が与えられる。情報は、「犯罪者aと犯罪者bは異なる組織に属する」という形式で与えられる。入力列の途中で「犯罪者aと犯罪者bは同じ組織に属するか?」という質問が与えられるので、それまでに得た情報に基づいて "In the same gang.", "In different gangs.", "Not sure yet." のいずれかを出力するプログラムをかけ。

解法

Union-Find木を使う。木は 2 * N 個のノードからなる。
0..(N-1) : 犯罪者i
N..(2N-1) : 番号jは犯罪者j-Nと違う組織に属するという意味
これによって、たとえば union(1, N+2) は、「犯罪者1は、犯罪者2と違う組織に属する」という意味になる。このようにして「違う組織に属する」という情報をうまく処理すれば、あとは単純。

#include <iostream>
#include <cmath>
#include <vector>
#include <queue>
#include <utility>
#include <algorithm>
using namespace std;

//------- Union-Find木 ---------
vector<int> par; // 親
vector<int> rnk; // 深さ
// 初期化
void init(int n) {
  par  = vector<int>(n);
  rnk = vector<int>(n);
  for (int i = 0; i < n; i++) {
    par[i] = i;
    rnk[i] = 0;
  }
}

// 木の根を求める
int find(int x) {
  if (par[x] == x) {
    return x;
  } else {
    return par[x] = find(par[x]);
  }
}

// xとyの属する集合を併合
void unite(int x, int y) {
  x = find(x);
  y = find(y);
  if (x == y) return;

  if (rnk[x] < rnk[y]) {
    par[x] = y;
  } else {
    par[y] = x;
    if (rnk[x] == rnk[y]) rnk[x]++;
  }
}

// xとyが同じ集合に属するか否か
bool same(int x, int y) {
  return find(x) == find(y);
}
//-------------------

int main() {
  ifstream cin("../test.txt");
  int testcase;
  cin >> testcase;
  while(testcase--) {
    int N, M, a, b;
    char c;
    cin >> N >> M;
    // Union-Find木の初期化
    // N番目以降はアンチ属性用
    init(2 * N);
    while (M--) {
      cin >> c;
      cin >> a >> b;
      if (c == 'A') {
        if (same(a, b)) {
          cout << "In the same gang." << endl;
        } else if (same(a, b+N)) {
          cout << "In different gangs." << endl;
        } else {
          cout << "Not sure yet" << endl;
        }
      } else if (c == 'D') {
        unite(a, b + N);
        unite(b, a + N);
      }
    }
  }
}

AOJ2170: Marked Ancestor

Marked Ancestor | Aizu Online Judge

問題概要

ノードがN個の木が与えられる。ノードには1からNの番号がつけられ、根は1である。木に対して二つの操作を考える。
M v: ノードvにマークする。
Q v: ノードvに最も近いマークされた先祖の番号を返す。
Q個の操作が与えられるので、操作Qで返される番号の総和を出力せよ。

解法

どこがUnion-Find木を使うんだろうと思った。
ちょっと考えると、これは木を「分離」していく問題だ。でも、Union-Find木は木を「合併」するのに長けているデータ構造だ。どうやって利用すればいいのだろうか。と、ここまで考えて解答を見た。
なるほど。
クエリを後ろから処理するらしい。マークを外していく。そうするとUnion-Find木の要領で合併ができるな。

#include <fstream>
#include <iostream>
#include <cmath>
#include <vector>
#include <map>
#include <queue>
#include <utility>
#include <algorithm>
using namespace std;

//------- Union-Find木 ---------
vector<int> par; // 親
// 初期化
void init(int n) {
  par = vector<int>(n);
  par[0] = 0;
}

// 木の根を求める
int find(int x) {
  if (par[x] == x) {
    return x;
  } else {
    return par[x] = find(par[x]);
  }
}

// xとyの属する集合を併合
void unite(int x, int y) {
  x = find(x);
  y = find(y);
  if (x == y) return;
  par[x] = y;
}

// xとyが同じ集合に属するか否か
bool same(int x, int y) {
  return find(x) == find(y);
}
//-------------------

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // ifstream cin ("2170_input.txt");
  int N, Q;
  while (1) {
    cin >> N >> Q;
    if (!(N | Q)) break;
    init(N);
    int a, i;
    for (i = 1; i < N; i++) {
      cin >> a;
      par[i] = a - 1;
    }
    vector<pair<char, int> > query(Q);
    char q;
    for (i = 0; i < Q; i++) {
      cin >> q >> a;
      query[i] = make_pair(q, a);
    }

    // 先にマークする。つまり木を分離する。親となるノードを覚えておく(ori_par)。
    // 同じノードに2回以上マークしている場合も考慮する(ori_par_cnt)。
    map<int, int> ori_par;
    map<int, int> ori_par_cnt;
    for (i = 0; i < Q; i++) {
      if (query[i].first == 'M') {
        a = query[i].second - 1;
        if (!ori_par_cnt.count(a)) {
          ori_par[a] = par[a];
          ori_par_cnt[a] = 1;
          par[a] = a;
        } else {
          ori_par_cnt[a]++;
        }
      }
    }

    // クエリを後ろから処理する
    long long sum = 0;
    reverse(query.begin(), query.end());
    for (i = 0; i < Q; i++) {
      a = query[i].second - 1;
      if (query[i].first == 'Q') {
        sum += find(a) + 1;
      } else if (query[i].first == 'M') {
        if (ori_par_cnt[a] == 1) {
          unite(a, ori_par[a]);
        } else {
          ori_par_cnt[a]--;
        }
      }
    }

    // 出力
    cout << sum << endl;
  }
}

感想

最近の方針は、10分くらい考えてアルゴリズムが思いつかなかったらググって調べる。長時間考えて答えが出せるようになるためには経験の蓄積が必要だと思っているので、まだ競技プログラミングの問題を解く絶対量が足りないから、わからなかったら思考に時間をかけずに調べることにしている。実装は自分でやるけど。短い時間で効果的に上達していきたいので、上達の効率を考えながら時間を投資していかないとね。