List Removals
Problem Overview
| Attribute | Value |
|---|---|
| CSES Link | List Removals |
| Difficulty | Medium |
| Category | Range Queries / Data Structures |
| Time Limit | 1 second |
| Key Technique | Segment Tree for Order Statistics |
Learning Goals
After solving this problem, you will be able to:
- Use a segment tree to maintain dynamic order statistics
- Perform binary search within a segment tree to find the k-th active element
- Handle dynamic deletion queries efficiently in O(log n) time
- Recognize problems that require “find k-th element” operations
Problem Statement
Problem: Given a list of n integers, process n queries where each query removes and returns the element at a specified position in the current list.
Input:
- Line 1: n (number of elements)
- Line 2: n space-separated integers (the initial list)
- Lines 3 to n+2: n integers, each representing a position p (1-indexed) to remove
Output:
- n lines: the element removed at each query
Constraints:
- 1 <= n <= 2 x 10^5
- 1 <= x_i <= 10^9
- 1 <= p <= current list size
Example
Input:
5
2 6 1 4 2
3
1
3
1
2
Output:
1
2
2
6
4
Explanation:
- Query p=3: List is [2,6,1,4,2], remove position 3 -> remove 1, list becomes [2,6,4,2]
- Query p=1: List is [2,6,4,2], remove position 1 -> remove 2, list becomes [6,4,2]
- Query p=3: List is [6,4,2], remove position 3 -> remove 2, list becomes [6,4]
- Query p=1: List is [6,4], remove position 1 -> remove 6, list becomes [4]
- Query p=2: List is [4], but wait - we only have 1 element? (Re-check: this was for illustration)
Let me use the simpler example from the problem:
Input:
5
1 2 3 4 5
3
3
2
1
Output:
3
2
1
- Query p=3: List [1,2,3,4,5], remove position 3 -> 3, list becomes [1,2,4,5]
- Query p=2: List [1,2,4,5], remove position 2 -> 2, list becomes [1,4,5]
- Query p=1: List [1,4,5], remove position 1 -> 1, list becomes [4,5]
Intuition: How to Think About This Problem
Pattern Recognition
Key Question: How do we efficiently find the k-th element in a dynamically shrinking list?
The naive approach of actually removing elements from an array takes O(n) per operation due to shifting. Instead, we can think of “removing” an element as marking it as deleted, and then use a data structure to quickly find the k-th non-deleted element.
Breaking Down the Problem
- What are we looking for? The k-th active (non-deleted) element at each query
- What information do we have? A list of values and their positions
- What’s the relationship between input and output? We need to map “logical position k” to “actual array index”
Key Insight
A segment tree can store the count of active elements in each range. To find the k-th element:
- If the left subtree has >= k active elements, the k-th element is in the left subtree
- Otherwise, it’s in the right subtree (and we search for position k - left_count)
This is essentially binary search within the segment tree, giving us O(log n) per query.
Solution 1: Brute Force (TLE)
Idea
Maintain the actual list and remove elements directly using Python’s list or vector operations.
Algorithm
- Read the initial list
- For each query position p:
- Remove and print the element at position p-1 (0-indexed)
Code
def solve_brute_force():
"""
Brute force: O(n) per removal due to shifting.
Time: O(n^2) total
Space: O(n)
"""
n = int(input())
arr = list(map(int, input().split()))
results = []
for _ in range(n):
p = int(input())
removed = arr.pop(p - 1) # O(n) operation
results.append(removed)
print('\n'.join(map(str, results)))
Complexity
| Metric | Value | Explanation |
|---|---|---|
| Time | O(n^2) | Each pop() shifts O(n) elements |
| Space | O(n) | Storing the array |
Why This Is Too Slow
With n = 2 x 10^5, we need O(n^2) = 4 x 10^10 operations, far exceeding the time limit.
Solution 2: Segment Tree for Order Statistics (Optimal)
Key Insight
The Trick: Use a segment tree where each node stores the count of active elements in its range. Finding the k-th element becomes a binary search within the tree.
Segment Tree Structure
| Node | Meaning |
|---|---|
tree[v] |
Count of active (non-deleted) elements in the range represented by node v |
Root: Contains total active elements (starts at n, decreases by 1 each query)
Algorithm
- Build: Initialize segment tree with all counts = 1 (all elements active)
- Query (find k-th):
- At each node, compare k with left child’s count
- Go left if k <= left_count, else go right with k = k - left_count
- When reaching a leaf, return that index
- Update: Set the found element’s count to 0 (mark deleted)
Dry Run Example
Input: n=5, arr=[1,2,3,4,5], queries: [3, 2, 1]
Initial Segment Tree (counts):
[5] <- root: 5 active elements
/ \
[2] [3] <- left subtree: 2, right: 3
/ \ / \
[1] [1] [1] [2]
| | | / \
idx0 idx1 idx2 [1] [1]
1 2 3 idx3 idx4
4 5
Query p=3: Find 3rd active element
At root: left_count = 2, k = 3
k > left_count, go RIGHT, k = 3 - 2 = 1
At right child [3]: left_count = 1, k = 1
k <= left_count, go LEFT
At leaf idx2: Found index 2, value = arr[2] = 3
Mark deleted: tree update, count at idx2 becomes 0
After Query 1:
[4] <- now 4 active
/ \
[2] [2] <- right is now 2
/ \ / \
[1] [1] [0] [2] <- idx2 is now 0
Query p=2: Find 2nd active element
At root: left_count = 2, k = 2
k <= left_count, go LEFT
At left child [2]: left_count = 1, k = 2
k > left_count, go RIGHT, k = 2 - 1 = 1
At leaf idx1: Found index 1, value = arr[1] = 2
Query p=1: Find 1st active element
At root: left_count = 1, k = 1
k <= left_count, go LEFT
At left child [1]: left_count = 1, k = 1
k <= left_count, go LEFT
At leaf idx0: Found index 0, value = arr[0] = 1
Results: [3, 2, 1]
Visual Diagram
Original: [1, 2, 3, 4, 5]
0 1 2 3 4 <- actual indices
Query p=3: Find 3rd active
[1, 2, 3, 4, 5]
1 2 * 4 5 <- position 3 is index 2
^
Remove arr[2]=3
After: [1, 2, _, 4, 5] (logically [1,2,4,5])
0 1 X 3 4 <- index 2 is "deleted"
Query p=2: Find 2nd active in [1,2,4,5]
1 2 _ 4 5
^
Remove arr[1]=2
Query p=1: Find 1st active in [1,4,5]
1 _ _ 4 5
^
Remove arr[0]=1
Code (Python)
import sys
from sys import stdin
input = stdin.readline
def solve():
"""
Segment Tree solution for List Removals.
Time: O(n log n)
Space: O(n)
"""
n = int(input())
arr = list(map(int, input().split()))
# Segment tree: stores count of active elements in each range
# Size 4*n is safe upper bound for segment tree array
tree = [0] * (4 * n)
def build(v, tl, tr):
"""Build segment tree. Each leaf starts with count 1."""
if tl == tr:
tree[v] = 1
else:
tm = (tl + tr) // 2
build(2 * v, tl, tm)
build(2 * v + 1, tm + 1, tr)
tree[v] = tree[2 * v] + tree[2 * v + 1]
def find_kth_and_remove(v, tl, tr, k):
"""
Find the k-th active element and mark it as deleted.
Returns the index of the k-th active element.
"""
# Decrease count as we will remove one element from this subtree
tree[v] -= 1
if tl == tr:
# Reached leaf node, this is the k-th element
return tl
tm = (tl + tr) // 2
left_count = tree[2 * v]
if k <= left_count:
# k-th element is in the left subtree
return find_kth_and_remove(2 * v, tl, tm, k)
else:
# k-th element is in the right subtree
return find_kth_and_remove(2 * v + 1, tm + 1, tr, k - left_count)
# Build the tree
build(1, 0, n - 1)
# Process queries
results = []
for _ in range(n):
p = int(input())
idx = find_kth_and_remove(1, 0, n - 1, p)
results.append(arr[idx])
print('\n'.join(map(str, results)))
if __name__ == "__main__":
solve()
Complexity
| Metric | Value | Explanation |
|---|---|---|
| Time | O(n log n) | n queries, each O(log n) tree traversal |
| Space | O(n) | Segment tree of size 4n |
Alternative: Binary Indexed Tree (Fenwick Tree)
A BIT can also solve this problem with binary search:
Note: The BIT approach uses external binary search, giving O(n log^2 n). The segment tree approach with built-in descent is faster at O(n log n), but both pass within time limits.
Common Mistakes
Mistake 1: Using 0-indexed positions
# WRONG: Treating p as 0-indexed
idx = find_kth_and_remove(1, 0, n - 1, p - 1) # Off by one!
# CORRECT: p is 1-indexed in the problem
idx = find_kth_and_remove(1, 0, n - 1, p)
Problem: The problem states positions are 1-indexed. Fix: Pass p directly without subtracting 1.
Mistake 2: Forgetting to decrement tree counts
# WRONG: Not updating counts during traversal
def find_kth(v, tl, tr, k):
if tl == tr:
return tl
# Missing: tree[v] -= 1
Problem: The tree counts become incorrect after removals.
Fix: Decrement tree[v] at the start of the function before recursing.
Mistake 3: Wrong tree size
Problem: Segment tree can have up to 4n nodes in worst case. Fix: Always allocate 4*n space for segment tree arrays.
Edge Cases
| Case | Input | Expected | Why It Matters |
|---|---|---|---|
| Single element | n=1, p=1 | Output that element | Simplest case |
| Remove first repeatedly | p=1 for all queries | Elements in order | Tests left-heavy traversal |
| Remove last repeatedly | p=current_size | Elements in reverse | Tests right-heavy traversal |
| Large values | arr[i] = 10^9 | Handle correctly | Values don’t affect logic |
| All same values | arr = [5,5,5,5,5] | Same output | Values are independent of positions |
When to Use This Pattern
Use Segment Tree for Order Statistics When:
- You need to find the k-th smallest/largest in a dynamic set
- Elements are being inserted or deleted
- You need O(log n) per operation
- The problem involves “position in sorted order” or “rank”
Don’t Use When:
- The array is static (use prefix sums or simpler structures)
- You only need minimum/maximum (simpler segment tree suffices)
- n is very small (brute force may be simpler)
Pattern Recognition Checklist:
- “Find element at position k” after deletions? -> Segment tree order statistics
- “How many elements less than x”? -> Segment tree / BIT
- Dynamic insertions AND deletions with rank queries? -> Consider balanced BST or segment tree
Related Problems
Easier (Do These First)
| Problem | Why It Helps |
|---|---|
| Static Range Sum Queries | Basic segment tree / prefix sum |
| Dynamic Range Sum Queries | Segment tree with point updates |
Similar Difficulty
| Problem | Key Difference |
|---|---|
| Josephus Problem I | Circular removal, similar technique |
| Josephus Problem II | Larger skip values, needs order statistics |
| Salary Queries | Coordinate compression + order statistics |
Harder (Do These After)
| Problem | New Concept |
|---|---|
| Range Updates and Sums | Lazy propagation |
| Polynomial Queries | Complex lazy propagation |
| Distinct Values Queries | Offline + segment tree |
Key Takeaways
- Core Idea: Use segment tree node values as counts to enable binary search within the tree
- Time Optimization: From O(n) per deletion to O(log n) by avoiding actual array shifts
- Space Trade-off: O(4n) extra space for segment tree
- Pattern: “Find k-th active element” = Segment Tree Order Statistics
Practice Checklist
Before moving on, make sure you can:
- Implement segment tree order statistics from scratch
- Explain why the binary descent within segment tree is O(log n)
- Recognize “dynamic k-th element” problems
- Implement both segment tree and BIT solutions
- Handle 1-indexed vs 0-indexed correctly