WeightedUnionFind

1
WeightedUnionFind<distance_class> wuf(num_elements);

1
uf.merge(a,b,w);

1
wuf.isSame(a,b);

1
wuf.size(a);

1
wuf.diff(a,b);

1
template <class T>
2
class WeightedUnionFind {
3
public:
4
  WeightedUnionFind(int n) : parents(n, -1), weights(n, 0) {}
5
  int find(int i) {
6
    if (parents[i] < 0)
7
      return i;
8
    weights[i] += weights[parents[i]];
9
    return parents[i] = find(parents[i]);
10
  }
11
  void merge(int a, int b, T w) {
12
    w += weight(a) - weight(b);
13
    a = find(a), b = find(b);
14
    if (a != b) {
15
      if (-parents[a] < -parents[b]) {
16
        swap(a, b);
17
        w = -w;
18
      }
19
      parents[a] += parents[b];
20
      parents[b] = a;
21
      weights[b] = w;
22
    }
23
  }
24
  T diff(int a, int b) {
25
    return weight(b) - weight(a);
26
  }
27
  bool connected(int a, int b) {
28
    return find(a) == find(b);
29
  }
30
  int size(int i) {
31
    return -parents[find(i)];
32
  }
33
private:
34
  vector<int> parents;
35
  vector<T> weights;
36
  T weight(int i) {
37
    find(i);
38
    return weights[i];
39
  }
40
};