CF 1249B1 - Books Exchange (easy version)

Rating: 1000
Tags: dsu, math
Solve time: 5m 26s
Verified: no

Solution

Problem Understanding

Each child starts with exactly one book, and every day all books are passed simultaneously according to a fixed permutation. If a child i gives their book to p[i], then after one day the book moves one step along this directed edge. After two days it moves two steps along the permutation, and so on.

From the perspective of a single starting book, we are repeatedly applying the same permutation. The question for each starting position i is: after how many applications of the permutation does the book that started at i return to i for the first time.

This is equivalent to asking for the length of the cycle that contains node i in a directed graph where each node has exactly one outgoing edge. Since p is a permutation, the graph is a disjoint union of cycles.

The constraints are small: n is at most 200 and there are at most 200 queries. This immediately tells us that even an O(n²) or O(n³) approach per test case is safe. We do not need heavy optimization or data structures beyond simple traversal and marking.

A subtle edge case is when a node maps to itself, meaning p[i] = i. In this case, the cycle length is 1 and the answer must be 1. Another case is a long cycle where no node repeats until the full cycle completes, for example a permutation like 2 3 4 5 1, where every node shares the same answer 5. Any incorrect solution that tries to simulate day-by-day movement separately for each node may still work here due to small constraints, but risks inefficiency or redundant recomputation.

Approaches

A direct simulation approach is to take each starting index i and repeatedly apply the permutation until we return to i, counting steps along the way. This is correct because each step exactly follows the rules of book passing. However, in the worst case, a single cycle of length n forces O(n) steps per starting node, leading to O(n²) per test case. With n up to 200 and q up to 200, this still passes comfortably (about 8 million operations worst-case), but it repeats work heavily since every node in a cycle recomputes the same traversal.

The key structural insight is that the permutation decomposes into disjoint cycles. Every node in the same cycle shares the same return time, equal to the cycle length. Instead of simulating separately, we traverse each cycle once, compute its size, and assign that size to all nodes in it. This reduces redundant work and turns the solution into a clean graph traversal problem.

We can find cycles by iterating over all nodes and starting a walk whenever we encounter an unvisited node. We follow p[i], p[p[i]], and so on until we return to the starting point, counting how many nodes are in that loop.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force Simulation	O(n²) per query	O(1) extra	Accepted
Cycle Decomposition	O(n) per query	O(n)	Accepted

Algorithm Walkthrough

We process each query independently and build the answer array.

Maintain a visited array of size n initialized to false, and an answer array initialized to zeros. The visited array ensures we do not process the same cycle multiple times.
Iterate over each node i from 1 to n. If i is already visited, skip it because it belongs to a cycle already processed.
When we find an unvisited node i, we start traversing the cycle beginning at i. We store the start node and move through the permutation repeatedly using p[current].
While traversing, we collect all nodes in the current cycle into a list and mark them as visited. We also count how many steps we take until we return to the starting node.
Once we close the cycle (we reach the start again), we assign the computed cycle length to every node in the collected list.

The reason we explicitly collect nodes rather than assigning on the fly is that we only know the cycle length after completing the traversal.

Why it works

Every node in a permutation belongs to exactly one directed cycle. The traversal from any unvisited node follows deterministic edges and must eventually return to the starting node because the graph is finite and each node has exactly one outgoing edge. This guarantees that each traversal discovers a complete cycle without branching or ambiguity. Since the answer is defined as the first return time, which is exactly the cycle length, assigning the cycle size to all nodes in that component is correct.

Python Solution

import sys
input = sys.stdin.readline

q = int(input())
for _ in range(q):
    n = int(input())
    p = [0] + list(map(int, input().split()))
    
    visited = [False] * (n + 1)
    ans = [0] * (n + 1)
    
    for i in range(1, n + 1):
        if visited[i]:
            continue
        
        cycle = []
        cur = i
        
        while not visited[cur]:
            visited[cur] = True
            cycle.append(cur)
            cur = p[cur]
        
        length = len(cycle)
        for v in cycle:
            ans[v] = length
    
    print(*ans[1:])

The solution works by explicitly decomposing the permutation into cycles. The 1-indexing simplifies direct mapping from input to graph nodes. Each node is visited exactly once, so the traversal cost is linear per query.

A common implementation pitfall is forgetting to mark nodes as visited before moving to the next node in the cycle, which can lead to infinite loops in incorrect implementations. Another subtle issue is mixing up the termination condition of the cycle walk; we rely on revisiting a previously visited node within the current traversal, which is safe because every node has exactly one outgoing edge.

Worked Examples

Example 1

Input:

n = 5
p = [1, 2, 3, 4, 5]

Step	Start	Visited nodes	Current node	Cycle collected	Cycle length
1	1	{}	1	[]	-
2	1	{1}	1	[1]	-
3	1	{1,2}	2	[1,2]	-
4	1	{1,2,3}	3	[1,2,3]	-
5	1	{1,2,3,4}	4	[1,2,3,4]	-
6	1	{1,2,3,4,5}	5	[1,2,3,4,5]	-
7	1	stop	1	full cycle	5

All nodes receive answer 5, which matches the fact that the entire permutation is one cycle.

This confirms that the algorithm correctly handles full-cycle permutations.

Example 2

Input:

n = 4
p = [2, 1, 4, 3]

Step	Start	Visited nodes	Current node	Cycle collected	Cycle length
1	1	{}	1	[]	-
2	1	{1}	1	[1]	-
3	1	{1,2}	2	[1,2]	2
4	3	{}	3	[]	-
5	3	{3}	3	[3]	-
6	3	{3,4}	4	[3,4]	2

This shows two independent cycles, each of length 2, and confirms that visited tracking prevents recomputation.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n) per query	Each node is visited exactly once during cycle decomposition
Space	O(n)	Arrays for visited tracking and output storage

With n ≤ 200 and q ≤ 200, the total work is at most 40,000 node visits, which is trivial within limits.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    import sys
    input = sys.stdin.readline

    q = int(input())
    out_lines = []
    for _ in range(q):
        n = int(input())
        p = [0] + list(map(int, input().split()))
        
        visited = [False] * (n + 1)
        ans = [0] * (n + 1)
        
        for i in range(1, n + 1):
            if visited[i]:
                continue
            cycle = []
            cur = i
            while not visited[cur]:
                visited[cur] = True
                cycle.append(cur)
                cur = p[cur]
            for v in cycle:
                ans[v] = len(cycle)
        out_lines.append(" ".join(map(str, ans[1:])))
    
    return "\n".join(out_lines)

# provided samples
assert run("""6
5
1 2 3 4 5
3
2 3 1
6
4 6 2 1 5 3
1
1
4
3 4 1 2
5
5 1 2 4 3
""") == """1 1 1 1 1
3 3 3
2 3 3 2 1 3
1
2 2 2 2
4 4 4 1 4"""

# all self-loops
assert run("""1
3
1 2 3
""") == "1 1 1"

# two cycles
assert run("""1
6
2 1 4 3 6 5
""") == "2 2 2 2 2 2"

# single big cycle
assert run("""1
4
2 3 4 1
""") == "4 4 4 4"

Test input	Expected output	What it validates
identity permutation	all 1s	self-loop correctness
two disjoint swaps	all 2s	multiple cycles
single cycle	all n	full traversal correctness

Edge Cases

Self-loops such as p[i] = i are handled naturally because the traversal immediately returns to the starting node after visiting it once, producing cycle length 1. The visited marking ensures we stop correctly without infinite looping.

Large cycles where all nodes are connected in one loop are also handled correctly because the traversal continues until we revisit the start node, guaranteeing full coverage of the cycle before assignment.