Overview

This problem requires identifying groups of students connected by friendships and, for each notification, determining the size (number of people) of the group to which the notified student belongs.

Analysis

Key Insight

The key point of this problem is that “students who are directly or indirectly connected through friendships belong to the same group.”

For example, if student 1 and student 2 are friends, and student 2 and student 3 are friends, then student 1 and student 3 also belong to the same group. This is a problem of finding connected components.

Example: N=5, friendships: (1,2), (2,3), (4,5)
Group 1: {1, 2, 3} (size 3)
Group 2: {4, 5} (size 2)

Issues with the Naive Approach

If we count the group size using BFS/DFS for each query, a single query takes up to \(O(N + M)\). With \(Q\) queries, the overall complexity becomes \(O(Q \times (N + M))\). When \(N, M, Q\) are each up to \(2 \times 10^5\), this results in TLE (Time Limit Exceeded).

Solution

Since friendships do not change, we can precompute all groups and their sizes, allowing each query to be answered in \(O(1)\).

Union-Find (Disjoint Set Union) is the ideal data structure for managing groups.

Algorithm

What is Union-Find?

Union-Find is a data structure that efficiently supports the following two operations: - Find(x): Find the representative (root) of the group to which element \(x\) belongs - Union(x, y): Merge the groups of elements \(x\) and \(y\)

Solution Steps

1. Initialization: Each student starts in their own group (size 1)
2. Union Processing: For every friendship \((U_j, V_j)\), execute Union(U_j, V_j) to merge groups
3. Query Processing: For student \(S_k\), use Find(S_k) to find the group representative and output that group’s size

Concrete Example

N=5, M=3
Friendships: (1,2), (2,3), (4,5)

Initial state: {1}, {2}, {3}, {4}, {5} (each of size 1)

Union(1,2) → {1,2}, {3}, {4}, {5}
Union(2,3) → {1,2,3}, {4}, {5}
Union(4,5) → {1,2,3}, {4,5}

Query S=2 → Group size of Find(2)'s representative = 3
Query S=4 → Group size of Find(4)'s representative = 2

Complexity

• Time Complexity: \(O((N + M) \cdot \alpha(N) + Q)\)

  • Union-Find initialization: \(O(N)\)
  • \(M\) Union operations: \(O(M \cdot \alpha(N))\) (\(\alpha\) is the inverse Ackermann function, effectively constant)
  • \(Q\) queries: \(O(Q \cdot \alpha(N))\)
  • Overall, practically \(O(N + M + Q)\)

• Space Complexity: \(O(N + M)\)

  • Adjacency list: \(O(N + M)\)
  • Union-Find arrays: \(O(N)\)

Implementation Notes

1. Path Compression: In the find function, recursively update the parent to speed up future lookups

   def find(x):
      if parent[x] != x:
          parent[x] = find(parent[x])  # Path compression
      return parent[x]
   

2. Union by Rank: Attach the shorter tree under the taller tree to keep the tree height small

3. Size Management: When merging, also add up the group sizes

   size[px] += size[py]  # Add size to the merge target
   

4. Output Optimization: Accumulate results in a list and output them all at once at the end to reduce I/O overhead

   Source Code

import sys
from collections import deque

def main():
    input = sys.stdin.readline
    
    N, M = map(int, input().split())
    
    # 隣接リストを構築
    adj = [[] for _ in range(N + 1)]
    for _ in range(M):
        U, V = map(int, input().split())
        adj[U].append(V)
        adj[V].append(U)
    
    # 各生徒が属するグループとそのサイズを求める
    # Union-Findを使用
    parent = list(range(N + 1))
    rank = [0] * (N + 1)
    size = [1] * (N + 1)
    
    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    
    def union(x, y):
        px, py = find(x), find(y)
        if px == py:
            return
        if rank[px] < rank[py]:
            px, py = py, px
        parent[py] = px
        size[px] += size[py]
        if rank[px] == rank[py]:
            rank[px] += 1
    
    # 全ての友達関係をUnion-Findで処理
    for i in range(1, N + 1):
        for j in adj[i]:
            if i < j:
                union(i, j)
    
    Q = int(input())
    results = []
    for _ in range(Q):
        S = int(input())
        # Sが属するグループのサイズを出力
        root = find(S)
        results.append(size[root])
    
    print('\n'.join(map(str, results)))

if __name__ == "__main__":
    main()

---

This editorial was generated by claude4.5opus.