src/alfred/data_structure/weighted-dsu.hpp

• View this file on GitHub
• Last update: 2025-03-17 21:57:32+08:00
• Include: #include "src/alfred/data_structure/weighted-dsu.hpp"

Verified with

• verify/verify-yosupo-ds/yosupo-weighted-dsu.test.cpp

Code

#ifndef AFDS_WEIGHTED_DSU
#define AFDS_WEIGHTED_DSU

#include <numeric>
#include <vector>

template <class T>
struct WeightedDSU {
    std::vector<int> fa;
    std::vector<T> w;
    WeightedDSU(int n) : fa(n + 1), w(n + 1) {
        std::iota(fa.begin(), fa.end(), 0);
    }
    inline int find(int x) {
        if (fa[x] == x) return x;
        int f = fa[x], f2 = find(f);
        return w[x] += w[f], fa[x] = f2;
    }
    inline bool same(int x, int y) {
        return find(x) == find(y);
    }
    // Given info: a[x] + v = a[y]
    // Returns true if this operation has no conflict, false otherwise.
    inline bool merge(int x, int y, T v) {
        int fx = find(x), fy = find(y);
        if (fx == fy) {
            return w[x] + v == w[y];
        }
        w[fy] = w[x] + v - w[y], fa[fy] = fx;
        return true;
    }
    inline T distance(int x, int y) {
        return find(x), find(y), w[y] - w[x];
    }
};

#endif // AFDS_WEIGHTED_DSU

#line 1 "src/alfred/data_structure/weighted-dsu.hpp"



#include <numeric>
#include <vector>

template <class T>
struct WeightedDSU {
    std::vector<int> fa;
    std::vector<T> w;
    WeightedDSU(int n) : fa(n + 1), w(n + 1) {
        std::iota(fa.begin(), fa.end(), 0);
    }
    inline int find(int x) {
        if (fa[x] == x) return x;
        int f = fa[x], f2 = find(f);
        return w[x] += w[f], fa[x] = f2;
    }
    inline bool same(int x, int y) {
        return find(x) == find(y);
    }
    // Given info: a[x] + v = a[y]
    // Returns true if this operation has no conflict, false otherwise.
    inline bool merge(int x, int y, T v) {
        int fx = find(x), fy = find(y);
        if (fx == fy) {
            return w[x] + v == w[y];
        }
        w[fy] = w[x] + v - w[y], fa[fy] = fx;
        return true;
    }
    inline T distance(int x, int y) {
        return find(x), find(y), w[y] - w[x];
    }
};

Back to top page