Overview

This problem asks you to find the size of “connected components” in graph theory. Consider a graph where students are vertices and friendships are edges. For each query, you need to answer how many students are in the connected component (group) that the specified student belongs to.

Analysis

Rephrasing the Problem

The condition “connected directly or indirectly through friendships” refers to connected components in graph theory. When a teacher contacts a certain student, the message reaches everyone in the group that student belongs to. Therefore, the value we need to find is “the number of vertices in the connected component that student \(S_k\) belongs to.”

Naive Approach and Its Limitations

If we traverse the graph (using BFS or DFS) for each query \(Q\) to count the number of people in the group, the worst-case time complexity is \(O(Q \times (N + M))\). Since the constraints have \(N, M, Q\) all up to \(2 \times 10^5\), this is too slow.

Efficient Solution

By identifying all groups in advance and precomputing the size of each group, we can answer each query in \(O(1)\). 1. Build the graph. 2. Starting from each unvisited student, perform BFS (Breadth-First Search) or DFS (Depth-First Search) to identify all students belonging to the same group. 3. Record the size of that group and associate it with every student in the group. 4. Repeat until all students have been explored.

Algorithm

1. Graph Construction

Build a graph with students as vertices and friendships as undirected edges. In Python, adjacency lists are typically created using a list of lists, but for memory efficiency and speed, the provided code uses a flat array representation similar to CSR (Compressed Sparse Row) format.

2. Connected Component Search (BFS)

Start BFS from each unvisited vertex. - Use a queue (list of vertices to explore) to visit connected vertices one after another. - The set of vertices visited in a single BFS constitutes one group. - Record the size of that group (len(component_queue)) and store this value in the answer array (res_sizes) for all students belonging to the group.

3. Answering Queries

For each query \(S_k\), output the precomputed value res_sizes[S_k].

Complexity

• Time Complexity: \(O(N + M + Q)\)
  • \(O(N + M)\) for graph construction
  • \(O(N + M)\) for BFS traversal of all vertices (since each vertex and each edge is visited a constant number of times)
  • \(O(Q)\) for answering \(Q\) queries
• Space Complexity: \(O(N + M)\)
  • \(O(N + M)\) for storing the adjacency list
  • \(O(N)\) for storing visited flags and group size information

Implementation Notes

• Fast I/O: When handling more than \(10^5\) inputs in Python, reading all input at once with sys.stdin.read().split() and writing output at once with sys.stdout.write() is significantly faster than calling input() repeatedly.

• Graph Representation: A list of lists like adj = [[] for _ in range(N)] is intuitive, but can have significant overhead in Python. The technique used in the provided code, combining a 1D array with offsets (offsets), is effective for saving memory and improving speed.

• Recursion Limit: When using DFS, you need to increase the recursion depth with sys.setrecursionlimit, but this is not a concern when using BFS.

  Source Code

import sys

def solve():
    # Fast I/O: Read all input at once and convert to a list of integers
    # This approach is generally faster in Python than multiple input() or sys.stdin.readline() calls.
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # Map all string tokens to integers
    data = list(map(int, input_data))
    
    N = data[0]
    M = data[1]
    
    # Build the graph using Compressed Sparse Row (CSR) format.
    # This representation is memory-efficient and avoids the overhead of a list of lists.
    
    # Step 1: Count the degree of each node (student)
    degree = [0] * (N + 1)
    for i in range(M):
        u = data[2 + 2 * i]
        v = data[3 + 2 * i]
        degree[u] += 1
        degree[v] += 1
        
    # Step 2: Calculate the starting offsets for each student's adjacency list in a flat array
    # offsets[u-1] will be the start index of student u's neighbors in the 'adj' array.
    offsets = [0] * (N + 1)
    for i in range(N):
        offsets[i + 1] = offsets[i] + degree[i + 1]
        
    # Step 3: Populate the flat adjacency array
    adj = [0] * offsets[N]
    # Use a temporary pointer array to keep track of the next insertion position for each student
    curr_pos = list(offsets)
    for i in range(M):
        u = data[2 + 2 * i]
        v = data[3 + 2 * i]
        # Store neighbor v in u's slice
        adj[curr_pos[u - 1]] = v
        curr_pos[u - 1] += 1
        # Store neighbor u in v's slice
        adj[curr_pos[v - 1]] = u
        curr_pos[v - 1] += 1
        
    # Step 4: Identify connected components and calculate their sizes using BFS
    # res_sizes[i] will store the total number of students in the group containing student i.
    res_sizes = [0] * (N + 1)
    visited = [False] * (N + 1)
    
    # Identify where the query section begins in the flattened input data
    q_ptr_start = 2 + 2 * M
    Q = data[q_ptr_start]
    
    for i in range(1, N + 1):
        if not visited[i]:
            # Perform BFS to find all students in the same connected component
            component_queue = [i]
            visited[i] = True
            head = 0
            while head < len(component_queue):
                u = component_queue[head]
                head += 1
                
                # Retrieve neighbors of student u using CSR offsets
                start_idx = offsets[u - 1]
                end_idx = offsets[u]
                for j in range(start_idx, end_idx):
                    neighbor = adj[j]
                    if not visited[neighbor]:
                        visited[neighbor] = True
                        component_queue.append(neighbor)
            
            # The length of 'component_queue' is the size of the current group
            group_size = len(component_queue)
            for student in component_queue:
                res_sizes[student] = group_size
                
    # Step 5: Answer each query based on the calculated group sizes
    output_results = []
    for i in range(Q):
        target_student = data[q_ptr_start + 1 + i]
        output_results.append(str(res_sizes[target_student]))
        
    # Output all answers joined by newlines
    sys.stdout.write('\n'.join(output_results) + '\n')

if __name__ == '__main__':
    solve()

---

This editorial was generated by gemini-3-flash-preview.