DataStructure/DSU/DSU_weighted.h

• View this file on GitHub
• Last update: 2023-01-07 01:46:12+08:00

Verified with

• DataStructure/test/aizu_dsl_1_b_dsu_weighted.test.cpp

Code

template<class S>
struct WeightedDSU {
    std::vector<int> lab;
    std::vector<S> w;  // relative to parent

    WeightedDSU(int n) : lab(n, -1), w(n) {}

    int getRoot(int u) {
        if (lab[u] < 0) return u;
        return getRoot(lab[u]);
    }

    int weight(int u) {
        if (lab[u] < 0) return w[u];
        return w[u] + weight(lab[u]);
    }

    // weight(t) = weight(s) + diff
    // returns false if contradicts
    bool merge(int s, int t, S diff) {
        // jump to root
        diff = diff + weight(s) - weight(t);
        s = getRoot(s);
        t = getRoot(t);
        if (s == t) return diff == 0;
        if (lab[s] > lab[t]) {
            std::swap(s, t);
            diff = -diff;
        }
        lab[s] += lab[t];
        lab[t] = s;
        w[t] = diff;
        return true;
    }
};

#line 1 "DataStructure/DSU/DSU_weighted.h"
template<class S>
struct WeightedDSU {
    std::vector<int> lab;
    std::vector<S> w;  // relative to parent

    WeightedDSU(int n) : lab(n, -1), w(n) {}

    int getRoot(int u) {
        if (lab[u] < 0) return u;
        return getRoot(lab[u]);
    }

    int weight(int u) {
        if (lab[u] < 0) return w[u];
        return w[u] + weight(lab[u]);
    }

    // weight(t) = weight(s) + diff
    // returns false if contradicts
    bool merge(int s, int t, S diff) {
        // jump to root
        diff = diff + weight(s) - weight(t);
        s = getRoot(s);
        t = getRoot(t);
        if (s == t) return diff == 0;
        if (lab[s] > lab[t]) {
            std::swap(s, t);
            diff = -diff;
        }
        lab[s] += lab[t];
        lab[t] = s;
        w[t] = diff;
        return true;
    }
};

Back to top page