公式

F - 友達の輪 / Circle of Friends 解説 by admin

Claude 4.6 Opus (Thinking)

Overview

This is a problem that requires efficiently processing group merging (union) and group size queries (find + size). It can be solved efficiently using Union-Find (disjoint set data structure).

Analysis

Naive Approach and Its Issues

One could consider managing which group each student belongs to using an array, and rewriting the group number of every member in one group upon each merge. However, this method takes \(O(N)\) in the worst case for a single merge, resulting in \(O(NQ)\) overall for \(Q\) instructions. When \(N, Q\) are at most \(2 \times 10^5\), this amounts to \(4 \times 10^{10}\) operations, which will not fit within the time limit.

Key Insight

This problem is a classic “dynamically manage grouping of elements” problem, and Union-Find (disjoint set data structure) is a perfect fit. Using Union-Find, merging and searching can be performed in nearly \(O(1)\).

Algorithm

Union-Find is a data structure that represents each group as a tree structure.

  • parent[x]: Records the parent of student \(x\). The root is the representative of the group.
  • rank[x]: Records an upper bound on the height of the tree (used for union by rank).
  • size[x]: When \(x\) is a root, records the number of people in that group.

Main Operations

  1. find(x): Returns the root of the group that student \(x\) belongs to. Performs path compression (path splitting) along the way to speed up future searches.
  2. union(x, y): Merges the groups of students \(x\) and \(y\). Performs union by rank (attaching the tree with smaller rank under the tree with larger rank) to suppress the growth of tree height. Updates size during merging.
  3. get_size(x): Finds the root using find(x) and returns the size of that root.

Concrete Example

Consider \(N = 5\) with the following instructions:

  1. 1 1 2 → Merge students 1 and 2 → Groups: {1,2}, {3}, {4}, {5}
  2. 1 3 4 → Merge students 3 and 4 → Groups: {1,2}, {3,4}, {5}
  3. 2 1 → Number of people in student 1’s group → Output 2
  4. 1 2 3 → Merge the groups of students 2 and 3 → Groups: {1,2,3,4}, {5}
  5. 2 4 → Number of people in student 4’s group → Output 4

Complexity

  • Time complexity: \(O((N + Q) \cdot \alpha(N))\)
    • \(\alpha(N)\) is the inverse Ackermann function, which is practically almost constant (at most \(5\)).
    • By combining path compression and union by rank, the total cost of \(Q\) operations is nearly \(O(N + Q)\).
  • Space complexity: \(O(N)\)
    • Three arrays parent, rank, and size are maintained, each with \(N+1\) elements.

Implementation Notes

  • Handling 1-indexed: Since student numbers range from \(1\) to \(N\), we initialize with UnionFind(N + 1) (size \(N+1\)), leaving index \(0\) unused.

  • Path compression implementation: self.parent[x] = self.parent[self.parent[x]] flattens the tree with each find call (a technique called path splitting).

  • Batch output processing: Instead of calling print for each type-2 instruction, results are accumulated in a list and output all at once with '\n'.join(out) at the end, reducing I/O overhead.

  • Fast input: sys.stdin.readline is assigned to input to speed up Python’s standard input.

    Source Code

import sys
input = sys.stdin.readline

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.size = [1] * n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        self.size[rx] += self.size[ry]
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1

    def get_size(self, x):
        return self.size[self.find(x)]

def main():
    N, Q = map(int, input().split())
    uf = UnionFind(N + 1)
    out = []
    for _ in range(Q):
        query = list(map(int, input().split()))
        if query[0] == 1:
            a, b = query[1], query[2]
            uf.union(a, b)
        else:
            x = query[1]
            out.append(str(uf.get_size(x)))
    print('\n'.join(out))

main()

This editorial was generated by claude4.6opus-thinking.

投稿日時:
最終更新: