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

  1. What are we looking for? The k-th active (non-deleted) element at each query
  2. What information do we have? A list of values and their positions
  3. 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

  1. Read the initial list
  2. 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

  1. Build: Initialize segment tree with all counts = 1 (all elements active)
  2. 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
  3. 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

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

  1. Core Idea: Use segment tree node values as counts to enable binary search within the tree
  2. Time Optimization: From O(n) per deletion to O(log n) by avoiding actual array shifts
  3. Space Trade-off: O(4n) extra space for segment tree
  4. 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

Additional Resources