AIOJ

Introduction

A disjoint-set union (DSU) is a data structure for managing the sets to which elements belong. It is implemented as a forest, where each tree represents a set and the nodes in that tree represent the elements of the set.

As its name suggests, a DSU supports two operations:

Union (Unite): merge the sets containing two elements (that is, merge the corresponding trees).
Find (Find): find the set containing an element (that is, find the root of the corresponding tree); this can be used to determine whether two elements belong to the same set.

After suitable modifications, a DSU can support deleting or moving individual elements, or maintaining edge weights on its trees. A persistent DSU can also be implemented using a segment tree with dynamic node creation Persistent DSU.

Warning

A DSU cannot separate a set with low complexity.

Initialization

At the beginning, each element lies in its own set, represented by a tree consisting only of its root. For convenience, we set each root's parent to itself.

Implementation

C++

struct dsu {
  vector<size_t> pa;

  explicit dsu(size_t size) : pa(size) { iota(pa.begin(), pa.end(), 0); }
};

Python

class Dsu:
    def __init__(self, size):
        self.pa = list(range(size))

Find

We need to move upward along the tree until we reach the root.

Implementation

C++

size_t dsu::find(size_t x) { return pa[x] == x ? x : find(pa[x]); }

Python

def find(self, x):
    return x if self.pa[x] == x else self.find(self.pa[x])

Path Compression

Every element visited during a find belongs to that set, so we can connect it directly to the root to speed up future finds.

Implementation

C++

size_t dsu::find(size_t x) { return pa[x] == x ? x : pa[x] = find(pa[x]); }

Python

def find(self, x):
    if self.pa[x] != x:
        self.pa[x] = self.find(self.pa[x])
    return self.pa[x]

Union

To merge two trees, we only need to connect the root of one tree to the root of the other.

Implementation

C++

void dsu::unite(size_t x, size_t y) { pa[find(x)] = find(y); }

Python

def unite(self, x, y):
    self.pa[self.find(x)] = self.find(y)

Union by Heuristic

During a union, the choice of which tree's root becomes the root of the new tree affects the complexity of future operations. We can attach the tree with fewer nodes or smaller depth to the other tree to avoid degeneration.

Detailed Complexity Discussion

Because the only operations we need to support are union and find, whenever we merge two sets, either set can be attached below the other and the result will still be correct. However, different ways of attaching them have different time complexities. Specifically, if we attach a set tree with both fewer nodes and smaller depth beneath a larger set tree, subsequent find operations will clearly take less time than under the other scheme (and this also yields a better worst-case time complexity).

Of course, we cannot always find a set that is smaller in both the number of nodes and the depth as described above. Since both node count and depth are easy to maintain, we often choose one of them as the ranking function. Regardless of which one we choose, the time complexity is $O (m\alpha(m,n))$ ; see the papers cited in References for the proof.

In actual algorithm-contest code, even without union by heuristic, the code can often finish within the time limit. Tarjan's paper¹ proves that using only path compression, without union by heuristic, has a worst-case time complexity of $O (m \log n)$ . Yao's paper² proves that using only path compression, without union by heuristic, still has an average-case time complexity of $O (m\alpha(m,n))$ .

If we use only union by heuristic and not path compression, the time complexity is $O(m\log n)$ . Since path compression can cause a large number of modifications in a single union, it is sometimes unsuitable. For example, persistent DSU and segment-tree divide and conquer + DSU generally use DSU with only union by heuristic.

Reference implementation using union by size (note that the initialization method must be adjusted):

Implementation

C++

struct dsu {
  vector<size_t> pa, size;

  explicit dsu(size_t size_) : pa(size_), size(size_, 1) {
    iota(pa.begin(), pa.end(), 0);
  }

  void unite(size_t x, size_t y) {
    x = find(x), y = find(y);
    if (x == y) return;
    if (size[x] < size[y]) swap(x, y);
    pa[y] = x;
    size[x] += size[y];
  }
};

Python

class Dsu:
    def __init__(self, size):
        self.pa = list(range(size))
        self.size = [1] * size

    def unite(self, x, y):
        x, y = self.find(x), self.find(y)
        if x == y:
            return
        if self.size[x] < self.size[y]:
            x, y = y, x
        self.pa[y] = x
        self.size[x] += self.size[y]

Reference Implementation

The complete implementation of a DSU with path compression and union by size is shown below:

Template Problem Luogu P3367 [Template] Disjoint Set Union Reference Implementation

C++

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

struct DSU {
  std::vector<size_t> pa, size;

  explicit DSU(size_t size_) : pa(size_), size(size_, 1) {
    std::iota(pa.begin(), pa.end(), 0);
  }

  size_t find(size_t x) { return pa[x] == x ? x : pa[x] = find(pa[x]); }

  void unite(size_t x, size_t y) {
    x = find(x), y = find(y);
    if (x == y) return;
    if (size[x] < size[y]) std::swap(x, y);
    pa[y] = x;
    size[x] += size[y];
  }
};

int main() {
  int n, m;
  std::cin >> n >> m;
  DSU dsu(n + 1);
  for (; m; --m) {
    int z, x, y;
    std::cin >> z >> x >> y;
    if (z == 1) {
      dsu.unite(x, y);
    } else {
      std::cout << (dsu.find(x) == dsu.find(y) ? 'Y' : 'N') << '\n';
    }
  }
  return 0;
}

Python

class Dsu:
    def __init__(self, size):
        self.pa = list(range(size))
        self.size = [1] * size

    def find(self, x):
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        return self.pa[x]

    def unite(self, x, y):
        x, y = self.find(x), self.find(y)
        if x == y:
            return
        if self.size[x] < self.size[y]:
            x, y = y, x
        self.pa[y] = x
        self.size[x] += self.size[y]


if __name__ == "__main__":
    n, m = map(int, input().split())
    dsu = Dsu(n + 1)
    for _ in range(m):
        z, x, y = map(int, input().split())
        if z == 1:
            dsu.unite(x, y)
        else:
            print("Y" if dsu.find(x) == dsu.find(y) else "N")

Complexity

After using both path compression and union by heuristic, the average time per DSU operation is only $O(\alpha(n))$ . Here, $\alpha$ is the inverse Ackermann function, which grows extremely slowly. In other words, the average running time of a single DSU operation can be regarded as a very small constant. A proof of the time complexity is given on this page.

Inverse Ackermann Function

The Ackermann function $A(m, n)$ is defined as follows:

$A(m, n) = \begin{cases}n+1&\text{if }m=0\\A(m-1,1)&\text{if }m>0\text{ and }n=0\\A(m-1,A(m,n-1))&\text{otherwise}\end{cases}$

The inverse Ackermann function $\alpha(n)$ is the inverse of the Ackermann function; that is, the largest integer $m$ such that $A(m, m) \leqslant n$ .

The space complexity of a DSU is clearly $O(n)$ .

Extended Operations

Starting from an ordinary DSU, we can make a series of modifications to support more operations or maintain more complex information.

DSU with Deletion

An ordinary DSU cannot support deletion because deleting a node inevitably deletes all nodes in the subtree rooted at it. To solve this problem, a DSU with deletion can use virtual nodes to ensure that every node storing actual data is always a leaf. To do this, create a virtual node for every data node during initialization, and set each data node's parent to that virtual node. Since each union of two sets only connects the roots of the two trees, only virtual nodes ever have children. This ensures that deleting a node does not accidentally delete any other node.

After deleting an individual node, note that a new virtual node must be created for that node as its parent; otherwise, subsequent union and deletion operations cannot be performed correctly.

Template Problem SPOJ JMFILTER - Junk-Mail Filter Reference Implementation

C++

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

struct DSU {
  size_t id;
  std::vector<size_t> pa, size;

  explicit DSU(size_t size_, size_t m)
      : id(size_ * 2), pa(size_ * 2 + m), size(size_ * 2 + m, 1) {
    // The first half of size is not actually used; it is only here to simplify index calculations
    std::iota(pa.begin(), pa.begin() + size_,
              size_);  // Make i point to virtual node i + size_
    std::iota(pa.begin() + size_, pa.end(), size_);  // Make every virtual node point to itself
  }

  size_t find(size_t x) { return pa[x] == x ? x : pa[x] = find(pa[x]); }

  void unite(size_t x, size_t y) {
    x = find(x), y = find(y);
    if (x == y) return;
    if (size[x] < size[y]) std::swap(x, y);
    pa[y] = x;
    size[x] += size[y];
  }

  void erase(size_t x) {
    size_t y = find(x);
    --size[y];
    pa[x] = id++;
  }
};

int main() {
  int n, m, case_id = 0;
  while ((std::cin >> n >> m), n) {
    DSU dsu(n, m);
    for (; m; --m) {
      char ch;
      std::cin >> ch;
      if (ch == 'M') {
        int x, y;
        std::cin >> x >> y;
        dsu.unite(x, y);
      } else if (ch == 'S') {
        int x;
        std::cin >> x;
        dsu.erase(x);
      }
    }
    int res = 0;
    for (int i = n; i < dsu.id; ++i) {
      if (dsu.size[i] && i == dsu.find(i)) {
        ++res;
      }
    }
    std::cout << "Case #" << (++case_id) << ": " << res << '\n';
  }
  return 0;
}

Python

# This code is for illustration only; it cannot pass the original problem because it times out
class Dsu:
    def __init__(self, size, m):
        self.id = size * 2
        # Make i point to virtual node i + size_; every virtual node points to itself
        self.pa = list(range(size, size * 2)) + list(range(size, size * 2 + m))
        # The first half of size is not actually used; it is only here to simplify index calculations
        self.size = [1] * (size * 2 + m)

    def find(self, x):
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        return self.pa[x]

    def unite(self, x, y):
        x, y = self.find(x), self.find(y)
        if x == y:
            return
        if self.size[x] < self.size[y]:
            x, y = y, x
        self.pa[y] = x
        self.size[x] += self.size[y]

    def erase(self, x):
        y = self.find(x)
        self.size[y] -= 1
        self.pa[x] = self.id
        self.id += 1


if __name__ == "__main__":
    case_id = 0
    while True:
        n, m = map(int, input().split())
        if not n:
            break
        dsu = Dsu(n, m)
        for _ in range(m):
            op = input().split()
            if op[0] == "M":
                x = int(op[1])
                y = int(op[2])
                dsu.unite(x, y)
            elif op[0] == "S":
                x = int(op[1])
                dsu.erase(x)
        res = 0
        for i in range(n, dsu.id):
            if dsu.size[i] and i == dsu.find(i):
                res += 1
        case_id += 1
        print(f"Case #{case_id}: {res}")
        input()

Similar methods can also be used to move a single element between sets. See the examples for implementation details.

Weighted DSU

We can also define weights on the edges of a DSU, together with an operation on those weights during path compression, to solve more problems. For example, for the classic NOI 2001 Food Chain problem, we can maintain the additive group modulo $3$ on the edge weights. Problems that maintain edge weights modulo a small modulus can also be solved by splitting each point in the DSU into multiple states. This technique in this special case is also called a "category DSU" or an "expanded-domain DSU". The examples later in the article illustrate these approaches.

To maintain edge weights in a DSU, the edge weights must be stored in the child nodes. Thus, each node stores the weight of the edge from it to its parent. The edge weight needs to be adjusted only when a node's parent changes. In general, this can happen during path compression or when merging two nodes. For example, if the edge weight is the distance between the current node and its parent, then during path compression, whenever the current node's parent is replaced by the root, the distance from the old parent to the root must be added to the edge weight stored at the current node. Similarly, when merging the sets containing two nodes, we need to calculate the weight of the new edge connecting their roots.

Template Problem Library Checker - Unionfind with Potential Reference Implementation

C++

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

constexpr int M = 998244353;

struct DSU {
  std::vector<size_t> pa, size, dist;

  explicit DSU(size_t size_) : pa(size_), size(size_, 1), dist(size_) {
    std::iota(pa.begin(), pa.end(), 0);
  }

  size_t find(size_t x) {
    if (pa[x] == x) return x;
    size_t y = find(pa[x]);
    (dist[x] += dist[pa[x]]) %= M;
    return pa[x] = y;
  }

  bool unite(size_t x, size_t y, int d) {
    find(x), find(y);
    (d += M - dist[y]) %= M;
    (d += dist[x]) %= M;
    x = pa[x], y = pa[y];
    if (x == y) return d == 0;
    if (size[x] < size[y]) {
      std::swap(x, y);
      d = (M - d) % M;
    }
    pa[y] = x;
    size[x] += size[y];
    dist[y] = d;
    return true;
  }

  int check(size_t x, size_t y) {
    find(x), find(y);
    if (pa[x] != pa[y]) return -1;
    return (dist[y] - dist[x] + M) % M;
  }
};

int main() {
  int n, m;
  std::cin >> n >> m;
  DSU dsu(n);
  for (; m; --m) {
    int op;
    std::cin >> op;
    if (op) {
      int u, v;
      std::cin >> u >> v;
      std::cout << dsu.check(u, v) << '\n';
    } else {
      int u, v, x;
      std::cin >> u >> v >> x;
      std::cout << dsu.unite(u, v, x) << '\n';
    }
  }
  return 0;
}

Python

M = 998244353


class DSU:
    def __init__(self, size: int):
        self.pa = list(range(size))
        self.size = [1] * size
        self.dist = [0] * size

    def find(self, x: int) -> int:
        if self.pa[x] == x:
            return x
        y = self.find(self.pa[x])
        self.dist[x] = (self.dist[x] + self.dist[self.pa[x]]) % M
        self.pa[x] = y
        return y

    def unite(self, x: int, y: int, d: int) -> bool:
        self.find(x)
        self.find(y)
        d = (d + M - self.dist[y]) % M
        d = (d + self.dist[x]) % M
        x, y = self.pa[x], self.pa[y]
        if x == y:
            return d == 0
        if self.size[x] < self.size[y]:
            x, y = y, x
            d = (M - d) % M
        self.pa[y] = x
        self.size[x] += self.size[y]
        self.dist[y] = d
        return True

    def check(self, x: int, y: int) -> int:
        self.find(x)
        self.find(y)
        if self.pa[x] != self.pa[y]:
            return -1
        return (self.dist[y] - self.dist[x] + M) % M


if __name__ == "__main__":
    n, m = map(int, input().split())
    dsu = DSU(n)
    for _ in range(m):
        op, *rest = map(int, input().split())
        if op:
            u, v = rest
            print(dsu.check(u, v))
        else:
            u, v, x = rest
            print(int(dsu.unite(u, v, x)))

Examples

In algorithm contests, most problems that directly test DSU require designing a special structure for the particular problem.

UVa11987 Almost Union-Find

Implement a data structure similar to a DSU that supports the following operations:

Merge the sets containing two elements.
Move a single element to the set containing another element.
Query the size of the set containing an element and the sum of its elements.

Solution

This problem makes operations 1 and 3 easy to handle; the difficulty lies in operation 2. Suppose we want to move element $x$ to the set containing element $y$ . In an ordinary DSU, directly setting the parent of $x$ to the root of the set containing $y$ does not work, because this would move all elements in the subtree rooted at $x$ as well. The solution is to ensure that $x$ has no children. Therefore, when building the DSU, create a virtual node for every element $x$ , namely $\tilde x$ , and make the parent of $x$ point to the corresponding virtual node $\tilde x$ . Then, when merging two sets, because a root is always connected to another root and all roots are virtual nodes, only virtual nodes can have children, while none of the nodes storing actual elements has children. Moving an element is much easier in this setting.

Reference Implementation

C++

#include <cassert>
#include <iostream>
#include <numeric>
#include <vector>

using namespace std;

struct dsu {
  vector<size_t> pa, size, sum;

  explicit dsu(size_t size_)
      : pa(size_ * 2), size(size_ * 2, 1), sum(size_ * 2) {
    // The first half of size and sum is not actually used; it is only here to simplify index calculations
    iota(pa.begin(), pa.begin() + size_, size_);
    iota(pa.begin() + size_, pa.end(), size_);
    iota(sum.begin() + size_, sum.end(), 0);
  }

  void unite(size_t x, size_t y) {
    x = find(x), y = find(y);
    if (x == y) return;
    if (size[x] < size[y]) swap(x, y);
    pa[y] = x;
    size[x] += size[y];
    sum[x] += sum[y];
  }

  void move(size_t x, size_t y) {
    auto fx = find(x), fy = find(y);
    if (fx == fy) return;
    pa[x] = fy;
    --size[fx], ++size[fy];
    sum[fx] -= x, sum[fy] += x;
  }

  size_t find(size_t x) { return pa[x] == x ? x : pa[x] = find(pa[x]); }
};

int main() {
  size_t n, m, op, x, y;
  while (cin >> n >> m) {
    dsu dsu(n + 1);  // The element range is 1..n
    while (m--) {
      cin >> op;
      switch (op) {
        case 1:
          cin >> x >> y;
          dsu.unite(x, y);
          break;
        case 2:
          cin >> x >> y;
          dsu.move(x, y);
          break;
        case 3:
          cin >> x;
          x = dsu.find(x);
          cout << dsu.size[x] << ' ' << dsu.sum[x] << '\n';
          break;
        default:
          assert(false);  // not reachable
      }
    }
  }
  return 0;
}

Python

class Dsu:
    def __init__(self, size):
        # The first half of size and sum is not actually used; it is only here to simplify index calculations
        self.pa = list(range(size, size * 2)) * 2
        self.size = [1] * size * 2
        self.sum = list(range(size)) * 2

    def unite(self, x, y):
        x, y = self.find(x), self.find(y)
        if x == y:
            return
        if self.size[x] < self.size[y]:
            x, y = y, x
        self.pa[y] = x
        self.size[x] += self.size[y]
        self.sum[x] += self.sum[y]

    def move(self, x, y):
        fx, fy = self.find(x), self.find(y)
        if fx == fy:
            return
        self.pa[x] = fy
        self.size[fx] -= 1
        self.size[fy] += 1
        self.sum[fx] -= x
        self.sum[fy] += x

    def find(self, x):
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        return self.pa[x]


if __name__ == "__main__":
    while True:
        try:
            n, m = map(int, input().split())
            dsu = Dsu(n + 1)  # The element range is 1..n
            for _ in range(m):
                op_x_y = list(map(int, input().split()))
                op = op_x_y[0]
                if op == 1:
                    dsu.unite(op_x_y[1], op_x_y[2])
                elif op == 2:
                    dsu.move(op_x_y[1], op_x_y[2])
                elif op == 3:
                    x = dsu.find(op_x_y[1])
                    print(dsu.size[x], dsu.sum[x])
        except EOFError:
            break

Luogu P2024 [NOI2011] Food Chain

The animal kingdom has three types of animals, $A,B,C$ , whose food chain forms an interesting cycle: $A$ eats $B$ , $B$ eats $C$ , and $C$ eats $A$ .

There are $N$ animals, numbered $1 \sim N$ . Each animal is one of $A,B,C$ , but we do not know which one.

Someone describes the food-chain relationships among these $N$ animals using two statements:

The first statement is 1 X Y, meaning that $X$ and $Y$ are the same kind.
The second statement is 2 X Y, meaning that $X$ eats $Y$ .

This person makes $K$ statements about the $N$ animals, one after another, using the two forms above. Some of the $K$ statements are true and some are false. A statement is false if it satisfies any of the following three conditions; otherwise, it is true.

The current statement conflicts with one or more previous true statements.
In the current statement, either $X$ or $Y$ is greater than $N$ .
The current statement says that $X$ eats $X$ .

Your task is to output the total number of false statements given $N$ and $K$ .

Solution 1

Consider using a weighted DSU to maintain the food-chain relationships. If $x$ and $y$ are the same kind, then $x\equiv y\pmod 3$ ; if $x$ eats $y$ , then $x - y \equiv 1 \pmod 3$ . This transforms the problem into the template problem discussed above.

More specifically, for each statement, except for the obvious false statements with $x>n$ or $y>n$ , we need to determine whether $x$ and $y$ are already connected. If they are connected, compute the distance between them modulo the modulus and compare it with the information asserted by the statement; otherwise, connect them according to the information provided by the statement. Apart from the obvious cases, a statement is false if and only if the two mentioned nodes are already connected and their corresponding distance conflicts with the information asserted by the statement.

Reference Implementation 1

C++

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

constexpr int M = 3;

struct DSU {
  std::vector<size_t> pa, size, dist;

  explicit DSU(size_t size_) : pa(size_), size(size_, 1), dist(size_) {
    std::iota(pa.begin(), pa.end(), 0);
  }

  size_t find(size_t x) {
    if (pa[x] == x) return x;
    size_t y = find(pa[x]);
    (dist[x] += dist[pa[x]]) %= M;
    return pa[x] = y;
  }

  bool unite(size_t x, size_t y, int d) {
    find(x), find(y);
    (d += M - dist[y]) %= M;
    (d += dist[x]) %= M;
    x = pa[x], y = pa[y];
    if (x == y) return d == 0;
    if (size[x] < size[y]) {
      std::swap(x, y);
      d = (M - d) % M;
    }
    pa[y] = x;
    size[x] += size[y];
    dist[y] = d;
    return true;
  }
};

int main() {
  int n, m;
  std::cin >> n >> m;
  DSU dsu(n + 1);
  int res = 0;
  for (; m; --m) {
    int op, x, y;
    std::cin >> op >> x >> y;
    if (x > n || y > n)
      ++res;
    else
      res += !dsu.unite(x, y, op == 1 ? 0 : 1);
  }
  std::cout << res << std::endl;
  return 0;
}

Python

M = 3


class DSU:
    def __init__(self, size: int):
        self.pa = list(range(size))
        self.size = [1] * size
        self.dist = [0] * size

    def find(self, x: int) -> int:
        if self.pa[x] == x:
            return x
        y = self.find(self.pa[x])
        self.dist[x] = (self.dist[x] + self.dist[self.pa[x]]) % M
        self.pa[x] = y
        return y

    def unite(self, x: int, y: int, d: int) -> bool:
        self.find(x)
        self.find(y)
        d = (d + M - self.dist[y]) % M
        d = (d + self.dist[x]) % M
        x, y = self.pa[x], self.pa[y]
        if x == y:
            return d == 0
        if self.size[x] < self.size[y]:
            x, y = y, x
            d = (M - d) % M
        self.pa[y] = x
        self.size[x] += self.size[y]
        self.dist[y] = d
        return True


if __name__ == "__main__":
    n, m = map(int, input().split())
    dsu = DSU(n + 1)
    res = 0
    for _ in range(m):
        op, x, y = map(int, input().split())
        if x > n or y > n:
            res += 1
        else:
            res += not dsu.unite(x, y, 0 if op == 1 else 1)
    print(res)

Solution 2

Split an organism $x$ into three states. In an implementation, we can treat the different states as different elements directly:

A state in the same set as $x$ means that the state and $x$ belong to the same species;
A state in the same set as $x+n$ can be eaten by $x$ ;
A state in the same set as $x+2n$ can eat $x$ .

Thus, for a statement:

1 x y is false if and only if:
1. $x>N$ or $y>N$ ;
2. $y$ is in the same set as one of $x+n$ and $x+2n$ .
2 x y is false if and only if:
1. $x>N$ or $y>N$ ;
2. $y$ is in the same set as one of $x$ and $x+2n$ .
If it is true, merge the corresponding states.

Reference Implementation 2

C++

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

struct DSU {
  std::vector<size_t> pa, size;

  explicit DSU(size_t size_) : pa(size_), size(size_, 1) {
    std::iota(pa.begin(), pa.end(), 0);
  }

  size_t find(size_t x) { return pa[x] == x ? x : pa[x] = find(pa[x]); }

  void unite(size_t x, size_t y) {
    x = find(x), y = find(y);
    if (x == y) return;
    if (size[x] < size[y]) std::swap(x, y);
    pa[y] = x;
    size[x] += size[y];
  }
};

int main() {
  int n, m;
  std::cin >> n >> m;
  DSU dsu(n * 3 + 1);
  int res = 0;
  for (; m; --m) {
    int op, x, y;
    std::cin >> op >> x >> y;
    if (x > n || y > n)
      ++res;
    else if (op == 1) {
      if (dsu.find(x) == dsu.find(y + n) ||
          dsu.find(x) == dsu.find(y + (n << 1))) {
        ++res;
      } else {
        dsu.unite(x, y);
        dsu.unite(x + n, y + n);
        dsu.unite(x + n * 2, y + n * 2);
      }
    } else {
      if (dsu.find(x) == dsu.find(y) || dsu.find(x) == dsu.find(y + n)) {
        ++res;
      } else {
        dsu.unite(x, y + n * 2);
        dsu.unite(x + n, y);
        dsu.unite(x + n * 2, y + n);
      }
    }
  }
  std::cout << res << std::endl;
  return 0;
}

Python

class Dsu:
    def __init__(self, size):
        self.pa = list(range(size))
        self.size = [1] * size

    def find(self, x):
        if self.pa[x] != x:
            self.pa[x] = self.find(self.pa[x])
        return self.pa[x]

    def unite(self, x, y):
        x, y = self.find(x), self.find(y)
        if x == y:
            return
        if self.size[x] < self.size[y]:
            x, y = y, x
        self.pa[y] = x
        self.size[x] += self.size[y]


if __name__ == "__main__":
    n, m = map(int, input().split())
    dsu = Dsu(n * 3 + 1)
    res = 0
    for _ in range(m):
        op, x, y = map(int, input().split())
        if x > n or y > n:
            res += 1
        elif op == 1:
            if dsu.find(x) == dsu.find(y + n) or dsu.find(x) == dsu.find(y + (n << 1)):
                res += 1
            else:
                dsu.unite(x, y)
                dsu.unite(x + n, y + n)
                dsu.unite(x + n * 2, y + n * 2)
        else:
            if dsu.find(x) == dsu.find(y) or dsu.find(x) == dsu.find(y + n):
                res += 1
            else:
                dsu.unite(x, y + n * 2)
                dsu.unite(x + n, y)
                dsu.unite(x + n * 2, y + n)
    print(res)

ABC396E Min of Restricted Sum

Given integers $N, M$ and integer sequences of length $M$ , $X=(X_1,X_2,\ldots,X_M)$ , $Y=(Y_1,Y_2,\ldots,Y_M)$ , and $Z=(Z_1,Z_2,\ldots,Z_M)$ . It is guaranteed that all elements of $X$ and $Y$ are in the range from $1$ to $N$ .

A nonnegative integer sequence of length $N$ , $A=(A_1,A_2,\ldots,A_N)$ , is called a good integer sequence if and only if it satisfies the following condition:

For every integer $i$ satisfying $1 \leq i \leq M$ , $A_{X_i} \oplus A_{Y_i} = Z_i$ , where $\oplus$ denotes the bitwise XOR operation.

Determine whether such a good integer sequence exists. If it does, find the good integer sequence that minimizes the element sum $\displaystyle \sum_{i=1}^N A_i$ , and output the sequence.

Solution

XOR is a "same" or "different" relation on each individual binary bit. Therefore, by splitting all binary bits of $A_i$ apart, the XOR relations can be maintained with a weighted DSU (or category DSU). Within the same connected component, elements correspond to the same bit position of different numbers in $A$ . When calculating the answer, the elements in one connected component usually fall into two groups, whose values should differ; assigning $0$ to the larger group and $1$ to the other minimizes the total weight.

Reference Implementation

C++

#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>

constexpr int M = 2;

struct DSU {
  std::vector<size_t> pa, size, dist;

  explicit DSU(size_t size_) : pa(size_), size(size_, 1), dist(size_) {
    std::iota(pa.begin(), pa.end(), 0);
  }

  size_t find(size_t x) {
    if (pa[x] == x) return x;
    size_t y = find(pa[x]);
    (dist[x] += dist[pa[x]]) %= M;
    return pa[x] = y;
  }

  bool unite(size_t x, size_t y, int d) {
    find(x), find(y);
    (d += M - dist[y]) %= M;
    (d += dist[x]) %= M;
    x = pa[x], y = pa[y];
    if (x == y) return d == 0;
    if (size[x] < size[y]) {
      std::swap(x, y);
      d = (M - d) % M;
    }
    pa[y] = x;
    size[x] += size[y];
    dist[y] = d;
    return true;
  }
};

int main() {
  int n, m;
  std::cin >> n >> m;
  DSU dsu((n + 1) << 5);
  for (; m; --m) {
    int x, y, z;
    std::cin >> x >> y >> z;
    for (int i = 0; i < 31; ++i) {
      if (!dsu.unite((x << 5) | i, (y << 5) | i, (z >> i) & 1)) {
        std::cout << -1 << std::endl;
        return 0;
      }
    }
  }
  std::vector<int> a(n + 1), cnt((n + 1) << 5);
  for (int i = 1; i < ((n + 1) << 5); ++i) {
    dsu.find(i);
    if (dsu.dist[i]) ++cnt[dsu.pa[i]];
  }
  for (int i = 1; i <= n; ++i) {
    for (int j = 0; j < 31; ++j) {
      int x = (i << 5) | j, y = dsu.pa[x];
      if ((cnt[y] > dsu.size[y] / 2) ^ dsu.dist[x]) {
        a[i] |= 1 << j;
      }
    }
  }
  for (int i = 1; i <= n; ++i) std::cout << a[i] << ' ';
  std::cout << std::endl;
  return 0;
}

Python

M = 2


class DSU:
    def __init__(self, size: int):
        self.pa = list(range(size))
        self.size = [1] * size
        self.dist = [0] * size

    def find(self, x: int) -> int:
        if self.pa[x] == x:
            return x
        y = self.find(self.pa[x])
        self.dist[x] = (self.dist[x] + self.dist[self.pa[x]]) % M
        self.pa[x] = y
        return y

    def unite(self, x: int, y: int, d: int) -> bool:
        self.find(x)
        self.find(y)
        d = (d + M - self.dist[y]) % M
        d = (d + self.dist[x]) % M
        x, y = self.pa[x], self.pa[y]
        if x == y:
            return d == 0
        if self.size[x] < self.size[y]:
            x, y = y, x
            d = (M - d) % M
        self.pa[y] = x
        self.size[x] += self.size[y]
        self.dist[y] = d
        return True


if __name__ == "__main__":
    n, m = map(int, input().split())
    dsu = DSU((n + 1) << 5)
    for _ in range(m):
        x, y, z = map(int, input().split())
        for i in range(31):
            if not dsu.unite((x << 5) | i, (y << 5) | i, (z >> i) & 1):
                print(-1)
                exit()

    a = [0] * (n + 1)
    cnt = [0] * ((n + 1) << 5)

    for i in range(1, (n + 1) << 5):
        dsu.find(i)
        if dsu.dist[i]:
            cnt[dsu.pa[i]] += 1

    for i in range(1, n + 1):
        for j in range(31):
            x = (i << 5) | j
            y = dsu.pa[x]
            if (cnt[y] > dsu.size[y] // 2) ^ dsu.dist[x]:
                a[i] |= 1 << j

    print(" ".join(map(str, a[1:])))

Exercises

Other applications

Kruskal's algorithm in minimum spanning trees and Tarjan's algorithm in lowest common ancestors are both based on disjoint-set union.

See applications of disjoint-set union for related topics.

References and further reading

Zhihu answer: Is there really a binary path-compression optimization in disjoint-set union?
Gabow, H. N., & Tarjan, R. E. (1985). A Linear-Time Algorithm for a Special Case of Disjoint Set Union. JOURNAL OF COMPUTER AND SYSTEM SCIENCES, 30, 209-221.PDF
CSDN: Extended-domain DSU and weighted DSU

Footnotes

Tarjan, R. E., & Van Leeuwen, J. (1984). Worst-case analysis of set union algorithms. Journal of the ACM (JACM), 31(2), 245-281.ResearchGate PDF ↩
Yao, A. C. (1985). On the expected performance of path compression algorithms.SIAM Journal on Computing, 14(1), 129-133. ↩