3935. Power Update After K-th Largest Insertion I 🔒

MediumSegment TreeArrayHash TableMathSortingHeap (Priority Queue)

LeetCode ↗

Problem Pattern Idea Example Code Complex Pitfalls

Problem Description

You are given an integer array nums and an integer p.

You are also given a 2D integer array queries, where each queries[i] = [val_i, k_i]. A special property holds: the difference between consecutive k_i values is always less than 10.

For each query, you need to perform the following three steps in order:

Insert val_i into nums.
Let x be the k_i-th largest element in the current nums.
Update p to p^x % (10^9 + 7), where p^x means p raised to the power of x.

Your task is to return an array ans, where ans[i] represents the value of p after processing the i-th query.

Key points to understand:

Each query is processed sequentially, and the modifications are persistent. That means after a query inserts val_i into nums, that element stays in nums for all subsequent queries.
Similarly, the value of p carries over from one query to the next. The p used in query i+1 is the updated p produced at the end of query i.
The k_i-th largest element means: if you sort the current nums in descending order, x is the element at position k_i (1-indexed). For example, the 1-st largest is the maximum element.
The hint that consecutive k_i values differ by less than 10 suggests that the position we are tracking shifts only slightly between queries, which allows efficient maintenance of the k_i-th largest element without recomputing from scratch each time.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

How We Pick the Algorithm

Why Ordered Set / Fenwick / Segment Tree?

This problem maps to Ordered Set / Fenwick / Segment Tree through a short path in the full flowchart.

Range queries with point updates require a segment tree for efficient answers.

Open in Flowchart

Intuition

The core challenge of this problem is repeatedly finding the k_i-th largest element after inserting new values into nums. A naive approach would sort the entire array for every query, but that would be far too slow when there are many queries.

Let's think about what we actually need. For each query, after inserting a value, we only care about one specific element: the k_i-th largest. So instead of keeping the whole array sorted and searching every time, we can be smarter about how we organize the data.

The trick is to split the elements into two groups:

A group r (right side) that holds the k_i largest elements.
A group l (left side) that holds everything else (the smaller elements).

If we keep this invariant — that r contains exactly the top k_i elements and all elements in l are smaller than those in r — then the answer x is simply the smallest element in r. Why? Because r holds the k_i largest elements, so its minimum is exactly the k_i-th largest of the whole array.

Now, why is this efficient? Here is where the hint about consecutive k_i values differing by less than 10 becomes important. Between two queries, the value of k barely changes (by less than 10). This means the boundary between l and r only shifts a tiny bit. So after each insertion, we only need to move a small number of elements across the boundary to restore the invariant, rather than rebuilding everything.

The maintenance logic for each query becomes:

Insert the new value into r first.
Move the smallest element of r over to l (this is just a balancing step to keep r from holding the newly inserted value if it doesn't belong there).
If r now has fewer than k elements, pull the largest elements from l back into r until r reaches size k.
If r now has more than k elements, push the smallest elements from r into l until r shrinks to size k.

Because k changes by less than 10 each time, steps 3 and 4 only loop a handful of times, making each query very fast.

Once the boundary is fixed, reading x = r[0] gives us the k_i-th largest element directly. Finally, we update p using fast exponentiation (pow(p, x, mod)), which computes p^x % (10^9 + 7) efficiently, and record the result. Carrying p forward into the next query naturally handles the chained updates described in the problem.

Pattern Learn more about Segment Tree, Math, Sorting and Heap (Priority Queue) patterns.

Solution Approach

Solution 1: Two Sorted Sets

We use two sorted sets, l and r, to maintain the current array nums. All elements in l are less than or equal to those in r, and the number of elements in r is kept equal to k_i. With this invariant, the smallest element of r is exactly the k_i-th largest element of the whole array.

Data structures used:

l — a SortedList (initially empty) that holds the smaller elements.
r — a SortedList initialized with all elements of nums, which will hold the largest k_i elements.
ans — the answer array.
mod = 10**9 + 7 — the modulus for the power operation.

Step-by-step walkthrough of each query [val, k]:

Insert the new value. We add val into r with r.add(val). Since r is a sorted list, the element is placed in its correct sorted position automatically.
Balance once. We move the smallest element of r into l using l.add(r.pop(0)). Here r.pop(0) removes and returns the smallest element of r (the one at index 0). This single move keeps the two sets balanced before we adjust to the exact target size.
Grow r if it is too small. If len(r) < k, we repeatedly pull the largest element from l back into r:
```
while len(r) < k:
    r.add(l.pop())
```
l.pop() (with no argument) removes and returns the largest element of l, which correctly belongs in the top-k group.
Shrink r if it is too big. If len(r) > k, we repeatedly push the smallest element of r down into l:
```
while len(r) > k:
    l.add(r.pop(0))
```
Only one of step 3 or step 4 will actually run, and because consecutive k values differ by less than 10, the loop iterates only a small number of times.
Read the answer. After balancing, r holds exactly the k largest elements, so x = r[0] (the smallest of r) is the k_i-th largest element of the current nums.
Update p with fast exponentiation. We compute p = pow(p, x, mod), which efficiently calculates p^x % (10^9 + 7) in O(log x) time. We then append the updated p to ans. Because p is reassigned, its new value carries over into the next query, matching the chained-update requirement.

Complexity analysis:

Let n be the initial size of nums and q be the number of queries. Inserting, popping, and indexing in a SortedList each cost O(log n). The balancing loops run only a constant number of times per query (thanks to the "less than 10" property). The power computation costs O(log x).
Time complexity: O((n + q) log(n + q) + q log x).
Space complexity: O(n + q) for storing all elements across the two sorted lists plus the answer array.

Example Walkthrough

Consider nums = [2, 5, 3], p = 2, and queries = [[7, 2], [4, 1]].

Initialization: l = [], r = [2, 3, 5] (all of nums), ans = [].

Query 1 [val=7, k=2]:

Insert: r = [2, 3, 5, 7]
Balance once: pop smallest 2 → l = [2], r = [3, 5, 7]
Grow? len(r)=3 < 2? No.
Shrink? len(r)=3 > 2? Yes → pop 3 → l = [2, 3], r = [5, 7]
Answer: x = r[0] = 5 (the 2nd largest of {2,3,5,7})
Update: p = pow(2, 5) = 32 → ans = [32]

Query 2 [val=4, k=1]: (p = 32)

Insert: r = [4, 5, 7]
Balance once: pop smallest 4 → l = [2, 3, 4], r = [5, 7]
Grow? len(r)=2 < 1? No.
Shrink? len(r)=2 > 1? Yes → pop 5 → l = [2, 3, 4, 5], r = [7]
Answer: x = r[0] = 7 (the 1st largest of {2,3,4,5,7})
Update: p = pow(32, 7) % (1e9+7) = 359738130 → ans = [32, 359738130]

Result: ans = [32, 359738130].

The boundary between l and r shifted only slightly between queries (k went 2 → 1), so the balancing loops ran just once each — exactly the efficiency the "k differs by less than 10" hint guarantees.

Solution Implementation

1from sortedcontainers import SortedList
2
3
4class Solution:
5    def powerUpdate(
6        self, nums: list[int], p: int, queries: list[list[int]]
7    ) -> list[int]:
8        # `left` holds the smaller elements; `right` holds the larger elements.
9        left = SortedList()
10        right = SortedList(nums)
11
12        answer: list[int] = []
13        MOD = 10**9 + 7
14
15        for value, k in queries:
16            # Insert the new value into the "right" (larger) bucket.
17            right.add(value)
18            # Move the smallest element of `right` into `left` to make room.
19            left.add(right.pop(0))
20
21            # Ensure `right` ends up containing exactly the k largest elements.
22            # If `right` is too small, pull the largest elements back from `left`.
23            while len(right) < k:
24                right.add(left.pop())
25            # If `right` is too large, push its smallest elements down to `left`.
26            while len(right) > k:
27                left.add(right.pop(0))
28
29            # `right[0]` is the smallest among the k largest, i.e. the k-th largest value.
30            kth_largest = right[0]
31
32            # Update the running power: p = p^(kth_largest) mod MOD.
33            p = pow(p, kth_largest, MOD)
34            answer.append(p)
35
36        return answer
37```
38
39**Notes / alternative perspectives:**
40
411. **Naming:** I renamed `l`/`r`/`val`/`x`/`mod` to descriptive names (`left`/`right`/`value`/`kth_largest`/`MOD`). The method name `powerUpdate` is kept unchanged per your requirement.
42
432. **Import:** `SortedList` comes from `sortedcontainers`, so I added the explicit import.
44
453. **Potential simplification:** Since each query adds exactly one element and `k` is given per query, you could alternatively use a single `SortedList` and index directly: after adding, `kth_largest = sl[-k]` (the k-th from the end). This avoids the two-bucket rebalancing entirely:
46
47```python3
48from sortedcontainers import SortedList
49
50
51class Solution:
52    def powerUpdate(
53        self, nums: list[int], p: int, queries: list[list[int]]
54    ) -> list[int]:
55        sl = SortedList(nums)
56        answer: list[int] = []
57        MOD = 10**9 + 7
58
59        for value, k in queries:
60            sl.add(value)              # insert the new value (O(log n))
61            kth_largest = sl[-k]       # k-th largest element (O(log n) indexing)
62            p = pow(p, kth_largest, MOD)
63            answer.append(p)
64
65        return answer
66

1class Solution {
2    // Insert value x into the multiset represented by TreeMap t (key -> count)
3    private void add(TreeMap<Integer, Integer> t, int x) {
4        // Increase the count of x by 1, defaulting to 0 if absent
5        t.merge(x, 1, Integer::sum);
6    }
7
8    // Remove a single occurrence of value x from the multiset t
9    private void remove(TreeMap<Integer, Integer> t, int x) {
10        int count = t.get(x);
11
12        if (count == 1) {
13            // Only one occurrence left, drop the key entirely
14            t.remove(x);
15        } else {
16            // Otherwise just decrement the count
17            t.put(x, count - 1);
18        }
19    }
20
21    // Fast exponentiation: compute (base^exp) % mod in O(log exp)
22    private long qpow(long base, int exp, int mod) {
23        long result = 1;
24
25        while (exp > 0) {
26            // If the lowest bit is set, multiply the accumulated result
27            if ((exp & 1) == 1) {
28                result = result * base % mod;
29            }
30
31            // Square the base and move to the next bit
32            base = base * base % mod;
33            exp >>= 1;
34        }
35
36        return result;
37    }
38
39    public List<Integer> powerUpdate(int[] nums, int p, int[][] queries) {
40        // Two multisets that together hold all current values.
41        // left  : the smaller portion (values below the boundary)
42        // right : the larger portion, whose smallest element is the answer source
43        TreeMap<Integer, Integer> left = new TreeMap<>();
44        TreeMap<Integer, Integer> right = new TreeMap<>();
45
46        // leftSize / rightSize track the number of elements in each multiset
47        int leftSize = 0;
48        int rightSize = nums.length;
49
50        // Initially all numbers go into the right multiset
51        for (int x : nums) {
52            add(right, x);
53        }
54
55        final int mod = 1_000_000_007;
56
57        List<Integer> ans = new ArrayList<>();
58
59        for (int[] query : queries) {
60            int val = query[0]; // value to insert for this query
61            int k = query[1];   // desired size of the right multiset
62
63            // Insert the new value into the right multiset
64            add(right, val);
65            ++rightSize;
66
67            // Move the current minimum from right to left so the boundary
68            // stays consistent before re-balancing
69            int v = right.firstKey();
70
71            remove(right, v);
72            --rightSize;
73
74            add(left, v);
75            ++leftSize;
76
77            // If the right side is too small, shift the largest elements
78            // from left back into right until it reaches size k
79            while (rightSize < k) {
80                v = left.lastKey();
81
82                remove(left, v);
83                --leftSize;
84
85                add(right, v);
86                ++rightSize;
87            }
88
89            // If the right side is too large, push its smallest elements
90            // over to left until it shrinks down to size k
91            while (rightSize > k) {
92                v = right.firstKey();
93
94                remove(right, v);
95                --rightSize;
96
97                add(left, v);
98                ++leftSize;
99            }
100
101            // The k-th relevant value is the smallest element of the right side
102            int x = right.firstKey();
103
104            // Chain the power: raise the running p to the power x modulo mod
105            p = (int) qpow(p, x, mod);
106
107            ans.add(p);
108        }
109
110        return ans;
111    }
112}
113

1class Solution {
2public:
3    using ll = long long;
4
5    // Insert a value into the multiset represented by the map (increment its count)
6    void add(map<int, int>& mp, int value) {
7        ++mp[value];
8    }
9
10    // Remove one occurrence of a value from the multiset; erase the key if count drops to zero
11    void remove(map<int, int>& mp, int value) {
12        if (--mp[value] == 0) {
13            mp.erase(value);
14        }
15    }
16
17    // Fast exponentiation: compute (base^exp) % mod
18    ll qpow(ll base, int exp, int mod) {
19        ll result = 1;
20
21        while (exp) {
22            // If the current lowest bit is set, multiply the result by the current base
23            if (exp & 1) {
24                result = result * base % mod;
25            }
26
27            // Square the base and shift the exponent right by one bit
28            base = base * base % mod;
29            exp >>= 1;
30        }
31
32        return result;
33    }
34
35    vector<int> powerUpdate(vector<int>& nums, int p, vector<vector<int>>& queries) {
36        // Two multisets used to partition values:
37        //   leftSet  -> smaller values (the "left" portion)
38        //   rightSet -> larger values (the "right" portion); its minimum is the boundary
39        map<int, int> leftSet, rightSet;
40
41        // leftSize / rightSize track the number of elements in each multiset
42        int leftSize = 0, rightSize = static_cast<int>(nums.size());
43
44        // Initially place all numbers in the right multiset
45        for (int x : nums) {
46            add(rightSet, x);
47        }
48
49        const int mod = 1e9 + 7;
50
51        vector<int> ans;
52
53        for (auto& q : queries) {
54            int val = q[0]; // value to insert for this query
55            int k   = q[1]; // target size of the right multiset
56
57            // Step 1: insert the new value into the right multiset
58            add(rightSet, val);
59            ++rightSize;
60
61            // Step 2: move the current smallest element from right to left
62            int boundary = rightSet.begin()->first;
63
64            remove(rightSet, boundary);
65            --rightSize;
66
67            add(leftSet, boundary);
68            ++leftSize;
69
70            // Step 3: if the right multiset is too small, pull the largest
71            // elements from the left multiset back into the right multiset
72            while (rightSize < k) {
73                int moved = leftSet.rbegin()->first;
74
75                remove(leftSet, moved);
76                --leftSize;
77
78                add(rightSet, moved);
79                ++rightSize;
80            }
81
82            // Step 4: if the right multiset is too large, push its smallest
83            // elements into the left multiset until the size matches k
84            while (rightSize > k) {
85                int moved = rightSet.begin()->first;
86
87                remove(rightSet, moved);
88                --rightSize;
89
90                add(leftSet, moved);
91                ++leftSize;
92            }
93
94            // The smallest value remaining in the right multiset is the exponent
95            int exponent = rightSet.begin()->first;
96
97            // Update p iteratively: p = p^exponent (mod 1e9 + 7)
98            p = static_cast<int>(qpow(p, exponent, mod));
99
100            ans.push_back(p);
101        }
102
103        return ans;
104    }
105};
106

1// A multiset backed by a Map (value -> count) plus a sorted array of distinct keys.
2// The sorted key array lets us access the minimum and maximum in O(log n) for updates.
3interface OrderedMultiset {
4    counts: Map<number, number>; // maps each value to its occurrence count
5    keys: number[];              // sorted list of distinct values currently present
6}
7
8// Create an empty ordered multiset.
9function createMultiset(): OrderedMultiset {
10    return { counts: new Map<number, number>(), keys: [] };
11}
12
13// Binary search: find the index of the leftmost key >= target in the sorted keys array.
14function lowerBound(keys: number[], target: number): number {
15    let low = 0;
16    let high = keys.length;
17
18    while (low < high) {
19        const mid = (low + high) >> 1;
20
21        if (keys[mid] < target) {
22            low = mid + 1;
23        } else {
24            high = mid;
25        }
26    }
27
28    return low;
29}
30
31// Insert a value into the multiset (increment its count); register the key if new.
32function add(mp: OrderedMultiset, value: number): void {
33    const current = mp.counts.get(value) ?? 0;
34    mp.counts.set(value, current + 1);
35
36    // Only insert into the sorted key array when this value first appears.
37    if (current === 0) {
38        const index = lowerBound(mp.keys, value);
39        mp.keys.splice(index, 0, value);
40    }
41}
42
43// Remove one occurrence of a value; erase the key entirely if its count drops to zero.
44function remove(mp: OrderedMultiset, value: number): void {
45    const current = mp.counts.get(value) ?? 0;
46
47    if (current - 1 === 0) {
48        mp.counts.delete(value);
49
50        // Drop the key from the sorted array when no occurrences remain.
51        const index = lowerBound(mp.keys, value);
52        mp.keys.splice(index, 1);
53    } else {
54        mp.counts.set(value, current - 1);
55    }
56}
57
58// Smallest value currently stored in the multiset.
59function getMin(mp: OrderedMultiset): number {
60    return mp.keys[0];
61}
62
63// Largest value currently stored in the multiset.
64function getMax(mp: OrderedMultiset): number {
65    return mp.keys[mp.keys.length - 1];
66}
67
68// Fast exponentiation: compute (base^exp) % mod using BigInt to avoid precision loss.
69function qpow(base: number, exp: number, mod: number): number {
70    let result = 1n;
71    let b = BigInt(base) % BigInt(mod);
72    const m = BigInt(mod);
73    let e = exp;
74
75    while (e > 0) {
76        // If the current lowest bit is set, multiply the result by the current base.
77        if (e & 1) {
78            result = (result * b) % m;
79        }
80
81        // Square the base and shift the exponent right by one bit.
82        b = (b * b) % m;
83        e >>= 1;
84    }
85
86    return Number(result);
87}
88
89function powerUpdate(nums: number[], p: number, queries: number[][]): number[] {
90    // Two multisets used to partition values:
91    //   leftSet  -> smaller values (the "left" portion)
92    //   rightSet -> larger values (the "right" portion); its minimum is the boundary
93    const leftSet = createMultiset();
94    const rightSet = createMultiset();
95
96    // leftSize / rightSize track the number of elements in each multiset.
97    let leftSize = 0;
98    let rightSize = nums.length;
99
100    // Initially place all numbers in the right multiset.
101    for (const x of nums) {
102        add(rightSet, x);
103    }
104
105    const mod = 1e9 + 7;
106
107    const ans: number[] = [];
108
109    for (const q of queries) {
110        const val = q[0]; // value to insert for this query
111        const k = q[1];   // target size of the right multiset
112
113        // Step 1: insert the new value into the right multiset.
114        add(rightSet, val);
115        ++rightSize;
116
117        // Step 2: move the current smallest element from right to left.
118        const boundary = getMin(rightSet);
119
120        remove(rightSet, boundary);
121        --rightSize;
122
123        add(leftSet, boundary);
124        ++leftSize;
125
126        // Step 3: if the right multiset is too small, pull the largest
127        // elements from the left multiset back into the right multiset.
128        while (rightSize < k) {
129            const moved = getMax(leftSet);
130
131            remove(leftSet, moved);
132            --leftSize;
133
134            add(rightSet, moved);
135            ++rightSize;
136        }
137
138        // Step 4: if the right multiset is too large, push its smallest
139        // elements into the left multiset until the size matches k.
140        while (rightSize > k) {
141            const moved = getMin(rightSet);
142
143            remove(rightSet, moved);
144            --rightSize;
145
146            add(leftSet, moved);
147            ++leftSize;
148        }
149
150        // The smallest value remaining in the right multiset is the exponent.
151        const exponent = getMin(rightSet);
152
153        // Update p iteratively: p = p^exponent (mod 1e9 + 7).
154        p = qpow(p, exponent, mod);
155
156        ans.push(p);
157    }
158
159    return ans;
160}
161

Time and Space Complexity

Time Complexity: O((n + m) log (n + m))

Let n be the length of nums and m be the length of queries.

Building the SortedList from nums takes O(n log n) time.
For each of the m queries, the following operations are performed:
- r.add(val) and l.add(...): insertions into a SortedList, each costing O(log (n + m)).
- r.pop(0) and l.pop(): removals from a SortedList, each costing O(log (n + m)).
- The two while loops rebalance the sizes of l and r. Although a single query might appear to do many moves, across all queries the total number of elements transferred between l and r is bounded. Each element moved costs O(log (n + m)), and the amortized number of moves per query is constant, so the total cost of all rebalancing across the m queries is O((n + m) log (n + m)).
- pow(p, x, mod): modular exponentiation costing O(log x), which is dominated by the other operations.
Therefore, the dominant term is O((n + m) log (n + m)).

Space Complexity: O(n + m)

The two SortedList structures l and r together hold all n original elements from nums plus up to m values added during the queries, giving O(n + m) space.
The answer list ans stores at most m results, contributing O(m).
Hence the overall space complexity is O(n + m).

Pattern Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Misunderstanding How `pow(p, x, MOD)` Chains Across Queries

The most common mistake is misreading the update rule and computing the power independently for each query instead of chaining it. The problem requires that the updated p from query i becomes the base for query i+1. A frequent buggy version looks like this:

# WRONG: uses the original p every time instead of the running value
for value, k in queries:
    sl.add(value)
    kth_largest = sl[-k]
    answer.append(pow(p, kth_largest, MOD))   # p is never reassigned!

Here p stays at its initial value forever, so every answer is p_initial ^ (kth_largest) rather than the cumulative result. Fix: reassign p each iteration so the new value carries forward:

p = pow(p, kth_largest, MOD)   # reassign, then append
answer.append(p)

Pitfall 2: Off-by-One Error When Indexing the k-th Largest

In the simplified single-SortedList solution, the k-th largest element is sl[-k], not sl[-k-1] or sl[k-1]. Because SortedList stores elements in ascending order:

sl[-1] is the largest (1st largest),
sl[-2] is the 2nd largest,
so sl[-k] is the k-th largest.

Writing sl[k - 1] accidentally returns the k-th smallest, which is a silent logic error that still runs without crashing. Always confirm whether your data structure is ascending or descending before indexing.

Pitfall 3: Forgetting That `r.pop(0)` Removes the Smallest, While `l.pop()` Removes the Largest

In the two-set approach, the rebalancing logic depends on which end you pop from:

right.pop(0) removes the smallest of right (index 0).
left.pop() (no argument) removes the last/largest element of left.

Mixing these up (e.g., calling right.pop() instead of right.pop(0)) breaks the invariant that all elements in left ≤ all elements in right, producing wrong answers. The invariant must hold after every balance step.

Pitfall 4: Applying the Modulus to the Exponent

A subtle but serious error is reducing the exponent x modulo MOD:

# WRONG: never take the exponent mod MOD
p = pow(p, kth_largest % MOD, MOD)

By Fermat's Little Theorem, exponents may only be reduced modulo MOD - 1 (and only when the base is coprime to MOD), never modulo MOD. Since the values here are array elements (small relative to MOD), no exponent reduction is needed at all — pass kth_largest directly. Reducing it incorrectly silently corrupts results.

Pitfall 5: Ignoring the Persistence of `nums`

Each query permanently inserts val_i. A common mistake is resetting or copying nums per query, treating insertions as temporary. This both gives wrong answers and wastes O(n) work per query. Maintain a single shared SortedList (or pair of sets) across all queries so that every inserted value remains for subsequent ones.

Pitfall 6: Over-Relying on the "Difference < 10" Hint

The hint about consecutive k values differing by less than 10 only guarantees the rebalancing loops run a small number of times in the two-set approach. It does not mean you can skip recomputing the k-th largest. If you switch to the single-SortedList solution, sl[-k] already handles any k in O(log n), so the hint becomes irrelevant — don't build fragile logic that assumes k barely changes.

Ready to land your dream job?

Unlock your dream job with a 5-minute quiz for a personalized study roadmap!

Get My Roadmap

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Get a Personalized Study Roadmap:

Which of the tree traversal order can be used to obtain elements in a binary search tree in sorted order?

Pre-order traversal

Post-order traversal

In-order traversal

Level-order traversal

I don't know