data-structure/weighted-union-find.hpp

• View this file on GitHub
• Last update: 2023-07-08 14:12:19+09:00
• Include: #include "data-structure/weighted-union-find.hpp"

Depends on

• data-structure/base.hpp

Code

#pragma once

#include "./base.hpp"

#include <numeric>
#include <vector>

namespace matumoto {
  template <typename T>
  struct WeightedUnionFind {
    vector<int> par;
    vector<int> siz;
    vector<T> w;
    WeightedUnionFind(int N): par(N), siz(N, 1), w(N, 0) {
      iota(par.begin(), par.end(), 0);
    }

    int size(int x) {
      return siz[x];
    }

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

    bool issame(int x, int y) {
      return root(x) == root(y);
    }

    void unite(int x, int y, T z) {
      z += weight(x);
      z -= weight(y);
      x = root(x);
      y = root(y);
      if (x == y)
        return;
      if (x < y) {
        swap(x, y);
        z = -z;
      }
      siz[x] += siz[y];
      par[y] = x;
      w[y] = z;
    }

    int root(int x) {
      if (x == par[x])
        return x;
      int rx = root(par[x]);
      w[x] += w[par[x]];
      par[x] = rx;
      return rx;
    }

    T weight(int x) {
      root(x);
      return w[x];
    }
  };
} // namespace matumoto

#line 2 "data-structure/weighted-union-find.hpp"

#line 2 "data-structure/base.hpp"

namespace matumoto {
  using namespace std;
}
#line 4 "data-structure/weighted-union-find.hpp"

#include <numeric>
#include <vector>

namespace matumoto {
  template <typename T>
  struct WeightedUnionFind {
    vector<int> par;
    vector<int> siz;
    vector<T> w;
    WeightedUnionFind(int N): par(N), siz(N, 1), w(N, 0) {
      iota(par.begin(), par.end(), 0);
    }

    int size(int x) {
      return siz[x];
    }

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

    bool issame(int x, int y) {
      return root(x) == root(y);
    }

    void unite(int x, int y, T z) {
      z += weight(x);
      z -= weight(y);
      x = root(x);
      y = root(y);
      if (x == y)
        return;
      if (x < y) {
        swap(x, y);
        z = -z;
      }
      siz[x] += siz[y];
      par[y] = x;
      w[y] = z;
    }

    int root(int x) {
      if (x == par[x])
        return x;
      int rx = root(par[x]);
      w[x] += w[par[x]];
      par[x] = rx;
      return rx;
    }

    T weight(int x) {
      root(x);
      return w[x];
    }
  };
} // namespace matumoto

Back to top page