CF 1858E2 - Rollbacks (Hard Version)

Rating: 2600
Tags: data structures, interactive, trees
Solve time: 1m 2s
Verified: yes

Solution

Problem Understanding

We are asked to maintain a dynamic array while processing four types of queries: append a number, remove the last $k$ numbers, roll back the last change (append or remove), and report the number of distinct numbers in the array. The array is initially empty. Every rollback undoes the most recent change, so multiple consecutive rollbacks can walk back through the history of changes.

The input is online: we receive one query at a time and must answer ? queries immediately before reading the next query. The number of queries $q$ can be up to $10^6$ and the total number of ? queries is up to $10^5$. Numbers appended can be as large as $10^6$.

Given the constraints, any approach that scans the array to count distinct elements on each ? query is too slow. A brute-force method could take $O(n^2)$ in the worst case, which is unacceptable with $q = 10^6$. Moreover, we need to support undo operations efficiently, which rules out simple use of Python lists or sets without a clever history-tracking mechanism.

Non-obvious edge cases include consecutive rollbacks, rolling back removals, and repeated elements. For example, consider this sequence:

+ 1
+ 2
- 2
!
?

After - 2, the array is empty. Rolling back the removal with ! restores [1,2]. A careless implementation that does not track full change history might fail to restore the array correctly or miscount distinct elements.

Approaches

A naive solution would maintain the array in a list, perform + and - operations, and count distinct elements by converting to a set for ? queries. Rollbacks could be handled by storing a full copy of the array after each change. This is correct, but copying the array repeatedly costs $O(n)$ per change. With $q = 10^6$, this is $O(n^2)$ time and memory, which is infeasible.

The key insight is that we do not need the full array copies to implement rollbacks. We can store a history of operations in a stack, where each operation records just enough information to undo it. For append + x, store ('add', x). For remove - k, store the list of removed elements. A rollback simply pops the last operation from this history and applies the inverse operation. To count distinct elements efficiently, we maintain a Counter of the current elements. This allows ? queries to run in $O(1)$ time since we can track the number of distinct elements incrementally. By combining a stack of changes with a frequency counter, we support append, remove, rollback, and distinct count efficiently.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(q * n)	O(n * q)	Too slow
Stack + Counter	O(q)	O(n)	Accepted

Algorithm Walkthrough

Initialize an empty list a to store the current array, a Counter object freq to track the frequency of each number, a variable distinct_count to store the current number of distinct elements, and a stack history to store past operations.
For each query, parse its type:
For + x, append x to a, increment freq[x]. If freq[x] was zero before, increment distinct_count. Push ('add', x) to history.
For - k, pop the last k elements from a and decrement their counts in freq. If any count drops to zero, decrement distinct_count. Push ('remove', [removed_elements]) to history.
For !, pop the last operation from history:

If it was ('add', x), remove the last element from a and decrement freq[x]. If freq[x] becomes zero, decrement distinct_count.
If it was ('remove', removed_elements), append removed_elements back to a and increment freq for each. If any count was zero before, increment distinct_count.

For ?, print distinct_count and flush output immediately.
Continue to the next query.

Why it works: The history stack guarantees that rollbacks always undo the last change correctly. The freq counter keeps track of how many times each number appears, ensuring that distinct_count is always accurate. All operations on a and freq are mirrored in history, so every rollback restores both a and the count of distinct numbers exactly.

Python Solution

import sys
from collections import Counter

input = sys.stdin.readline
output = sys.stdout.write

a = []
freq = Counter()
distinct_count = 0
history = []

q = int(input())
for _ in range(q):
    line = input()
    if line[0] == '+':
        x = int(line[2:])
        a.append(x)
        if freq[x] == 0:
            distinct_count += 1
        freq[x] += 1
        history.append(('add', x))
    elif line[0] == '-':
        k = int(line[2:])
        removed = []
        for _ in range(k):
            val = a.pop()
            freq[val] -= 1
            if freq[val] == 0:
                distinct_count -= 1
            removed.append(val)
        removed.reverse()  # preserve original order
        history.append(('remove', removed))
    elif line[0] == '!':
        op = history.pop()
        if op[0] == 'add':
            val = a.pop()
            freq[val] -= 1
            if freq[val] == 0:
                distinct_count -= 1
        else:
            removed = op[1]
            for val in removed:
                if freq[val] == 0:
                    distinct_count += 1
                freq[val] += 1
            a.extend(removed)
    else:  # '?'
        output(f"{distinct_count}\n")
        sys.stdout.flush()

The solution uses Counter to maintain frequencies efficiently. Each + or - operation updates distinct_count in $O(1)$ per element. Rollbacks mirror these updates. Flushing output after each ? ensures compliance with online query requirements.

Worked Examples

Sample 1 Trace

Query	Array `a`	History	freq	distinct_count	Output
+ 1	[1]	[('add',1)]	{1:1}	1	-
+ 2	[1,2]	[('add',1),('add',2)]	{1:1,2:1}	2	-
+ 2	[1,2,2]	[('add',1),('add',2),('add',2)]	{1:1,2:2}	2	-
?	[1,2,2]	same	same	2	2
!	[1,2]	[('add',1),('add',2)]	{1:1,2:1}	2	-
+ 3	[1,2,3]	[('add',1),('add',2),('add',3)]	{1:1,2:1,3:1}	3	-
- 2	[1]	[('add',1),('add',2),('add',3),('remove',[2,3])]	{1:1,2:0,3:0}	1	-
?	[1]	same	same	1	1
+ 1	[1,1]	...	{1:2,2:0,3:0}	1	-
?	[1,1]	...	same	1	1

This confirms distinct_count is updated correctly after appends, removals, and rollbacks.

Sample 2 Trace

Query	Array `a`	freq	distinct_count	Output
+1	[1]	{1:1}	1	-
+1000000	[1,1000000]	{1:1,1000000:1}	2	-
?	[1,1000000]	same	2	2
!	[1]	{1:1,1000000:0}	1	-
!	[]	{1:0,1000000:0}	0	-
?	[]	same	0	0

This confirms the algorithm handles rollbacks of + correctly and updates distinct counts even when elements disappear entirely.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(q)	Each query is processed in O(k) for removals, but total number of elements added/removed is ≤ q