Skip to content

Weighted DSU

When "Same Set" Isn't Enough

Potential function: compute relative difference between any two nodes

Weighted DSU with relative weights along parent pointers

Standard DSU answers one question: are A and B connected? But many problems need more. "What is the ratio of A to B?" "Are A and B in the same group or opposite groups?" "What is A's distance from B?" Plain connectivity can't answer these. You need a DSU that carries data along every edge.

Weighted DSU attaches a value to each parent link. When you find the root, you accumulate these values along the path. The result tells you the relationship between any node and its root -- and by extension, between any two nodes in the same component.

The Weight Array

In standard DSU, par[x] points to the parent. In weighted DSU, wt[x] stores the relationship between x and par[x].

The interpretation depends on the problem:

Problem type wt[x] means Operation
Ratios (LC 399) val(x) / val(par[x]) Multiply along path
Distances dist(x) - dist(par[x]) Add along path
Parity x and par[x] in same/different group XOR along path

Find with Weight Compression

Path compression in standard DSU just reassigns parent pointers. In weighted DSU, you must also update the weight to reflect the direct relationship to the new root.

int par[200005], rnk[200005];
double wt[200005]; // wt[x] = val(x) / val(par[x])

int find(int x) {
    if (par[x] == x) return x;
    int root = find(par[x]);
    wt[x] *= wt[par[x]]; // accumulate ratio to root
    par[x] = root;
    return root;
}

Before compression: x -> A -> B -> root, with wt[x] = val(x)/val(A), wt[A] = val(A)/val(B), wt[B] = val(B)/val(root).

After compression: x -> root, with wt[x] = val(x)/val(root) = wt[x] * wt[A] * wt[B].

💡 Key insight. The find function does two things at once: finds the root AND computes the cumulative weight. After find(x), wt[x] gives you the direct relationship between x and its root.

Unite with Weights

When uniting a and b with a known relationship w = val(a)/val(b), we need to compute the correct weight for the edge between their roots.

bool unite(int a, int b, double w) {
    int ra = find(a), rb = find(b);
    if (ra == rb) {
        // Already connected — check consistency
        return abs(wt[a] / wt[b] - w) < 1e-9;
    }
    // After find: wt[a] = val(a)/val(ra), wt[b] = val(b)/val(rb)
    // We want: val(ra)/val(rb) = ?
    // val(a)/val(b) = w
    // val(a) = wt[a] * val(ra), val(b) = wt[b] * val(rb)
    // So: wt[a]*val(ra) / (wt[b]*val(rb)) = w
    // val(ra)/val(rb) = w * wt[b] / wt[a]
    if (rnk[ra] < rnk[rb]) {
        par[ra] = rb;
        wt[ra] = w * wt[b] / wt[a];
    } else {
        par[rb] = ra;
        wt[rb] = wt[a] / (w * wt[b]);
        if (rnk[ra] == rnk[rb]) rnk[ra]++;
    }
    return true;
}

⚠ Gotcha. Call find(a) and find(b) before computing the root weight. The find calls update wt[a] and wt[b] to be relative to their respective roots. If you skip this, the algebra breaks.

Querying the Relationship

To get val(a)/val(b) for two nodes in the same component:

double query(int a, int b) {
    if (find(a) != find(b)) return -1.0; // undefined
    return wt[a] / wt[b];
}

After find(a), wt[a] = val(a)/val(root). After find(b), wt[b] = val(b)/val(root). So wt[a]/wt[b] = val(a)/val(b). Clean.

Full Trace: Ratio Queries

Variables: a, b, c, d. Process: a/b = 2.0, b/c = 3.0, query a/c, query a/d.

Operation Action par[] wt[] Notes
init -- [a,b,c,d] [1,1,1,1] Each node is its own root
a/b = 2.0 unite(a,b,2.0) [a,a,c,d] [1,0.5,1,1] par[b]=a, wt[b]=wt[a]/(2.0*wt[b])=0.5
b/c = 3.0 unite(b,c,3.0): find(b)->wt[b]=0.5, find(c)->wt[c]=1 [a,a,a,d] [1,0.5,0.167,1] par[c]=a, wt[c]=wt[b]/(3.0*wt[c])=0.167
query a/c find(a)->wt[a]=1, find(c)->wt[c]=0.167 -- -- wt[a]/wt[c] = 1/0.167 = 6.0
query a/d find(a) != find(d) -- -- Return -1.0 (not connected)

Result: a/c = 6.0 (correct: a/b * b/c = 2 * 3 = 6). a/d = -1.0 (unknown).

Parity DSU: The XOR Variant

Some problems ask: "Are A and B in the same group or opposite groups?" This is a parity relationship. Instead of multiplying weights, XOR them.

int par[200005], rnk[200005], rel[200005];
// rel[x]: 0 = same group as par[x], 1 = different group

int find(int x) {
    if (par[x] == x) return x;
    int root = find(par[x]);
    rel[x] ^= rel[par[x]];
    par[x] = root;
    return root;
}

// unite a and b, where r = 0 means same group, 1 means different
bool unite(int a, int b, int r) {
    int ra = find(a), rb = find(b);
    if (ra == rb) {
        return (rel[a] ^ rel[b]) == r; // consistency check
    }
    if (rnk[ra] < rnk[rb]) swap(ra, rb), swap(a, b);
    par[rb] = ra;
    rel[rb] = rel[a] ^ rel[b] ^ r;
    if (rnk[ra] == rnk[rb]) rnk[ra]++;
    return true;
}

After find(a) and find(b), rel[a] ^ rel[b] tells you whether a and b are in the same group (0) or different groups (1).

💡 Key insight. Parity DSU is just weighted DSU where the "weight" is a single bit and the "accumulation" is XOR. The structure is identical.

LC 399: Evaluate Division

Problem: Given equations like a/b = 2.0 and b/c = 3.0, answer queries like "what is a/c?"

This is a direct application of weighted DSU with multiplicative weights.

class Solution {
public:
    vector<double> calcEquation(
        vector<vector<string>>& equations,
        vector<double>& values,
        vector<vector<string>>& queries
    ) {
        unordered_map<string, int> id;
        int cnt = 0;
        for (auto& eq : equations)
            for (auto& s : eq)
                if (!id.count(s)) id[s] = cnt++;

        vector<int> par(cnt), rnk(cnt, 0);
        vector<double> wt(cnt, 1.0);
        iota(par.begin(), par.end(), 0);

        function<int(int)> find = [&](int x) -> int {
            if (par[x] == x) return x;
            int root = find(par[x]);
            wt[x] *= wt[par[x]];
            par[x] = root;
            return root;
        };

        auto unite = [&](int a, int b, double w) {
            int ra = find(a), rb = find(b);
            if (ra == rb) return;
            if (rnk[ra] < rnk[rb]) swap(ra, rb), swap(a, b), w = 1.0 / w;
            par[rb] = ra;
            wt[rb] = wt[a] / (w * wt[b]);
            if (rnk[ra] == rnk[rb]) rnk[ra]++;
        };

        for (int i = 0; i < (int)equations.size(); i++)
            unite(id[equations[i][0]], id[equations[i][1]], values[i]);

        vector<double> ans;
        for (auto& q : queries) {
            if (!id.count(q[0]) || !id.count(q[1])) {
                ans.push_back(-1.0);
            } else {
                int a = id[q[0]], b = id[q[1]];
                if (find(a) != find(b)) {
                    ans.push_back(-1.0);
                } else {
                    ans.push_back(wt[a] / wt[b]);
                }
            }
        }
        return ans;
    }
};

AtCoder arc111_b: Reversible Cards

Problem: N cards, each with a front value and back value. You can flip any card. Maximize the number of distinct values visible.

DSU insight: Each card creates an edge between its front and back values. Within each connected component of k nodes and e edges, you can choose at most min(k, e) distinct values to show. If the component has a cycle (e >= k), you can show all k values. If it's a tree (e = k-1), you can show k-1.

// For each component: if it has a cycle, contribute size.
// If it's a tree, contribute size - 1.

This uses standard DSU augmented with cycle detection: if find(a) == find(b) when processing an edge, that component has a cycle.

Try It

The starter code provides the weighted DSU skeleton. Implement find with weight compression and unite with weight algebra.

Test case 1: 

Equations: a/b = 2.0, b/c = 3.0
Query: a/c
Expected: 6.0

Test case 2: 

Equations: a/b = 2.0
Query: b/a
Expected: 0.5

Test case 3: 

Equations: a/b = 2.0
Query: a/x
Expected: -1.0 (x is not connected)

Challenge: Implement a parity DSU and solve the following: given pairs (a,b) labeled "same" or "different", determine if a contradiction exists.

Problems

Problem Link Difficulty
LC 399 — Evaluate Division leetcode.com/problems/evaluate-division Medium
LC 990 — Satisfiability of Equality Equations leetcode.com/problems/satisfiability-of-equality-equations Medium
AtCoder arc111_b — Reversible Cards atcoder.jp/contests/arc111/tasks/arc111_b Medium