structure/union-find/weighted-union-find.hpp

• View this file on GitHub
• Last update: 2024-05-01 23:44:47+09:00
• Include: #include "structure/union-find/weighted-union-find.hpp"

Verified with

• test/verify/aoj-dsl-1-b.test.cpp

Code

template <typename T>
struct WeightedUnionFind {
  vector<int> data;
  vector<T> ws;

  WeightedUnionFind() {}

  WeightedUnionFind(int sz) : data(sz, -1), ws(sz) {}

  int find(int k) {
    if (data[k] < 0) return k;
    auto par = find(data[k]);
    ws[k] += ws[data[k]];
    return data[k] = par;
  }

  T weight(int t) {
    find(t);
    return ws[t];
  }

  bool unite(int x, int y, T w) {
    w += weight(x);
    w -= weight(y);
    x = find(x), y = find(y);
    if (x == y) return false;
    if (data[x] > data[y]) {
      swap(x, y);
      w *= -1;
    }
    data[x] += data[y];
    data[y] = x;
    ws[y] = w;
    return true;
  }

  T diff(int x, int y) { return weight(y) - weight(x); }
};

#line 1 "structure/union-find/weighted-union-find.hpp"
template <typename T>
struct WeightedUnionFind {
  vector<int> data;
  vector<T> ws;

  WeightedUnionFind() {}

  WeightedUnionFind(int sz) : data(sz, -1), ws(sz) {}

  int find(int k) {
    if (data[k] < 0) return k;
    auto par = find(data[k]);
    ws[k] += ws[data[k]];
    return data[k] = par;
  }

  T weight(int t) {
    find(t);
    return ws[t];
  }

  bool unite(int x, int y, T w) {
    w += weight(x);
    w -= weight(y);
    x = find(x), y = find(y);
    if (x == y) return false;
    if (data[x] > data[y]) {
      swap(x, y);
      w *= -1;
    }
    data[x] += data[y];
    data[y] = x;
    ws[y] = w;
    return true;
  }

  T diff(int x, int y) { return weight(y) - weight(x); }
};

Back to top page