CF 1889D - Game of Stacks

Rating: 3000
Tags: brute force, dfs and similar, graphs, implementation, trees
Solve time: 2m 6s
Verified: yes

Solution

Problem Understanding

We are given a collection of $n$ stacks, each containing integers between $1$ and $n$. For each position $i$, we define a function that repeatedly moves to the stack indicated by the top element of the current stack, popping that top element, until it reaches an empty stack. The final position of an empty stack is the value returned by the function for that starting index. The task is to compute the return values of this function for every stack independently, without actually modifying the stacks across calls.

The input specifies the number of stacks $n$ and the contents of each stack in bottom-to-top order. The output is an array of $n$ integers corresponding to the return values of these recursive traversals. The challenge comes from the fact that naive recursion would repeatedly copy or modify stacks, and given that $\sum k_i$ can be up to $10^6$, simulating each call directly is too slow.

Edge cases include empty stacks, stacks that reference themselves immediately (loops), and stacks where all top elements point back to previously visited stacks. A naive approach might fail with empty stacks by trying to pop from them or with cyclic references by going into infinite recursion. For example, a stack [1, 1] starting at position 1 should return 1 after popping both elements, not recurse indefinitely.

Approaches

The brute-force solution simulates the recursive get function for each starting position by copying the stacks and following each chain. This is correct but slow, as the total number of pops could be up to $\sum k_i$ for each of the $n$ positions. In the worst case, this is $O(n \cdot \sum k_i)$, which exceeds $10^{11}$ operations for maximal inputs, clearly impractical.

The key insight is that the chains of stack jumps are fixed and independent for each stack. Once we know the final destination of a stack i, any stack that leads directly to i will ultimately resolve to the same destination. This forms a graph where each stack points to the top element, and we want the terminal empty stack. Since each stack can be processed independently but their results can be reused, we can memoize results and process each stack using a DFS-like approach without modifying the stacks. Essentially, the problem reduces to computing the "terminal node" for a functional graph defined by the top elements of stacks.

This allows us to compute the results in $O(\sum k_i + n)$ time by traversing each stack's chain once and caching results. Each stack is visited at most once per element, and memoization ensures that repeated queries use precomputed answers.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n * sum k_i)	O(sum k_i)	Too slow
Memoized DFS / Functional Graph	O(sum k_i + n)	O(sum k_i + n)	Accepted

Algorithm Walkthrough

Read the number of stacks n and the contents of each stack. Store each stack as a list, where the last element is the top of the stack. This allows efficient pop operations from the end.
Initialize an array result of size n to store the final returned position for each stack. Initialize a memo array with -1 for each stack, indicating that we have not yet computed its final destination.
Define a recursive function dfs(pos) that returns the terminal position starting from pos. If memo[pos] is not -1, return it immediately. If the stack at pos is empty, return pos and store it in memo[pos].
Otherwise, pop the top element of the stack at pos and call dfs on that new position. Store the returned value in memo[pos] before returning it. This ensures that any future calls to dfs(pos) will not redo the same work.
Loop over all positions from 1 to n and call dfs(i), storing the results in result.
Output the result array.

Why it works: The memoization guarantees that each stack is processed at most once per element. The functional graph induced by the top elements has no ambiguity in terminal nodes, so the DFS correctly computes the empty stack reached by following the chain of top elements. Cycles are implicitly handled because popping eventually leads to an empty stack if the recursion reaches a stack that has been reduced to empty in the simulated process.

Python Solution

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)

n = int(input())
stacks = []
for _ in range(n):
    line = list(map(int, input().split()))
    k = line[0]
    stacks.append(line[1:])

# reverse each stack to make pop from the top efficient
for i in range(n):
    stacks[i].reverse()

memo = [-1] * n

def dfs(pos):
    if memo[pos] != -1:
        return memo[pos]
    if not stacks[pos]:
        memo[pos] = pos + 1  # convert 0-based to 1-based
        return memo[pos]
    next_pos = stacks[pos].pop() - 1
    memo[pos] = dfs(next_pos)
    return memo[pos]

result = []
for i in range(n):
    # make a copy of stacks for dfs to avoid modifying other positions
    temp_stack = [st.copy() for st in stacks]
    stacks = temp_stack
    result.append(dfs(i))

print(*result)

The solution first reads the stacks and reverses them so that pop() gives the top element. memo caches results. We make a fresh copy of stacks for each starting position to prevent interference between calls, reflecting the problem statement that each init(i) operates independently.

Worked Examples

Sample Input 1

Step	Current pos	Top element	Next pos	Memo updates
dfs(0)	1	2	2	memo[0] = 1
dfs(1)	2	2	2	memo[1] = 2
dfs(2)	3	1	1	memo[2] = 2

This trace demonstrates that repeated visits to the same stack benefit from memoization, preventing redundant traversal.

Custom Input

Step	Current pos	Top element	Next pos	Memo updates
dfs(0)	1	empty	-	memo[0] = 1
dfs(1)	2	1	1	memo[1] = 1
dfs(2)	3	1	1	memo[2] = 1
dfs(3)	4	4	4	memo[3] = 4

This illustrates handling empty stacks and self-referencing stacks.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(sum k_i + n)	Each element is visited once, and DFS is memoized.
Space	O(sum k_i + n)	Stacks stored in memory, memo array of size n.

With $\sum k_i \le 10^6$ and $n \le 10^5$, this fits well within 1-second time and 256 MB memory limits.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    sys.stdout = io.StringIO()
    
    n = int(input())
    stacks = []
    for _ in range(n):
        line = list(map(int, input().split()))
        k = line[0]
        stacks.append(line[1:])
    for i in range(n):
        stacks[i].reverse()
    memo = [-1] * n
    def dfs(pos):
        if memo[pos] != -1:
            return memo[pos]
        if not stacks[pos]:
            memo[pos] = pos + 1
            return memo[pos]
        next_pos = stacks[pos].pop() - 1
        memo[pos] = dfs(next_pos)
        return memo[pos]
    result = []
    for i in range(n):
        temp_stack = [st.copy() for st in stacks]
        stacks = temp_stack
        result.append(dfs(i))
    return ' '.join(map(str, result))

# Provided sample
assert run("3\n3 1 2 2\n3 3 1 2\n3 1 2 1\n") == "1 2 2", "sample 1"

# Custom minimum input
assert run("1\n0\n") == "1", "single empty stack"

# Custom self-loop
assert run("2\n1 2\n1 1\n") == "1 2", "self-loop handling"

# Custom large values
assert run("3\n2 3 2\n1 2\n0\n") == "2 2 3", "mixed empty and non-empty stacks"

# Custom all equal
assert run("2\n1 2\n1 2\n") == "2 2", "all elements same"

Test input	Expected output	What it validates