TopCoder SRM 687 Div1 Medium: AllGraphCuts

問題

https://community.topcoder.com/stat?c=problem_statement&pm=14214

解法

Gomory-Hu 木が作れることからGomory-Hu 木が作れることは明らか。
kyuridenamida先生のコードを写経した。

コード

#include <bits/stdc++.h>
using namespace std;

struct UnionFind {
  vector<int> data;
  UnionFind(int size) : data(size, -1) {}
  UnionFind(const UnionFind &u) { data = u.data; }
  bool unionSet(int x, int y) {
    x = root(x);
    y = root(y);
    if (x != y) {
      if (data[y] < data[x]) swap(x, y);
      data[x] += data[y];
      data[y] = x;
    }
    return x != y;
  }
  bool findSet(int x, int y) { return root(x) == root(y); }
  int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]); }
  int size(int x) { return -data[root(x)]; }
  bool is_same(int x, int y) { return findSet(x, y); }
};

class AllGraphCuts {
 public:
  vector<int> findGraph(vector<int> x) {
    int N = sqrt(x.size()) + 0.1;

    // check
    for (int i = 0; i < N; ++i)
      if (x[N * i + i] != 0) return {-1};

    map<int, vector<int>> es;
    set<int> w;
    for (int pos = 0; pos < N * N; ++pos) {
      int weight = x[pos];
      es[weight].push_back(pos);
      w.insert(weight);
    }

    UnionFind uf(N);
    vector<int> ans;
    for (auto rit = w.rbegin(); rit != w.rend(); ++rit) {
      int weight = *rit;
      for (auto pos : es[weight]) {
        int i = pos / N;
        int j = pos % N;
        if (i != j && uf.is_same(i, j)) return {-1};
      }

      for (auto pos : es[weight]) {
        int i = pos / N;
        int j = pos % N;
        if (!uf.is_same(i, j)) {
          ans.push_back(weight * N * N + i * N + j);
          uf.unionSet(i, j);
        }
      }
    }

    return ans;
  }
};