Skip to content

3935. Power Update After K-th Largest Insertion I πŸ”’

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] = [vali, ki] and the difference between consecutive ki values is always less than 10.

For each query:

  • Insert vali into nums.
  • Let x be the kith largest element in the current nums.
  • Update p to px % (109 + 7).

Return an array ans where the ans[i] represents the value of p after processing the ith query.

 

Example 1:

Input: nums = [2], p = 4, queries = [[3,1],[1,2]]

Output: [64,4096]

Explanation:

i vali Current
nums		 ki kith
largest		 p New p = pk % (109 + 7)
0 3 [2, 3] 1 3 4 43 % (109 + 7) = 64
1 1 [2, 3, 1] 2 2 64 642 % (109 + 7) = 4096

Thus, ans = [64, 4096].

Example 2:

Input: nums = [7,5], p = 6, queries = [[4,3],[7,2]]

Output: [1296,220296870]

Explanation:

i vali Current​​​​​​​
nums		 ki kith
largest		 p New p = pk % (109 + 7)
0 4 [7, 5, 4] 3 4 6 64 % (109 + 7) = 1296
1 7 [7, 5, 4, 7] 2 7 1296 12967 % (109 + 7) = 220296870

Thus, ans = [1296, 220296870]

 

Constraints:

  • 1 <= nums.length <= 2 Γ— 104
  • 1 <= nums[i] <= 106
  • ​​​​​​​1 <= p <= 106
  • 1 <= queries.length <= 2 Γ— 104
  • ​​​​​​​1 <= vali <= 106
  • 1 <= ki <= n + i + 1
  • |ki - ki - 1| < 10 for i > 0

Solutions

Solution 1: Two Sorted Sets

Thinking

Adjacent queries change \(k\) by less than \(10\), so the side that holds the current \(k\)-th largest can be maintained as a sliding window. Two sorted sets \(l\) and \(r\) split the values so that \(r\) always stores the \(k\) largest.

A new value enters \(r\), then elements move between the two sides until \(|r|=k\); \(r[0]\) is the \(k\)-th largest and \(p\) is updated by modular exponentiation.

Because \(k\) barely moves, the amortized number of transfers stays small.

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 equal to \(k_i\).

For each query, we insert \(val_i\) into \(r\), then move the smallest element in \(r\) to \(l\) until the size of \(r\) becomes \(k_i\). At this point, the smallest element in \(r\) is the \(k_i\)-th largest element in the current \(nums\). We then use fast exponentiation to update \(p\) as \(p^x \bmod (10^9 + 7)\), and append the updated \(p\) to the answer array.

The time complexity is \(O((n + m) \log (n + m))\), and the space complexity is \(O(n + m)\), where \(n\) and \(m\) are the lengths of \(\textit{nums}\) and \(\textit{queries}\), respectively.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def powerUpdate(
        self, nums: list[int], p: int, queries: list[list[int]]
    ) -> list[int]:
        l = SortedList()
        r = SortedList(nums)
        ans = []
        mod = 10**9 + 7
        for val, k in queries:
            r.add(val)
            l.add(r.pop(0))
            while len(r) < k:
                r.add(l.pop())
            while len(r) > k:
                l.add(r.pop(0))
            x = r[0]
            p = pow(p, x, mod)
            ans.append(p)
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class Solution {
    private void add(TreeMap<Integer, Integer> t, int x) {
        t.merge(x, 1, Integer::sum);
    }

    private void remove(TreeMap<Integer, Integer> t, int x) {
        int v = t.get(x);

        if (v == 1) {
            t.remove(x);
        } else {
            t.put(x, v - 1);
        }
    }

    private long qpow(long a, int b, int mod) {
        long ans = 1;

        while (b > 0) {
            if ((b & 1) == 1) {
                ans = ans * a % mod;
            }

            a = a * a % mod;
            b >>= 1;
        }

        return ans;
    }

    public List<Integer> powerUpdate(int[] nums, int p, int[][] queries) {
        TreeMap<Integer, Integer> l = new TreeMap<>();
        TreeMap<Integer, Integer> r = new TreeMap<>();

        int sz1 = 0, sz2 = nums.length;

        for (int x : nums) {
            add(r, x);
        }

        int mod = 1_000_000_007;

        List<Integer> ans = new ArrayList<>();

        for (int[] q : queries) {
            int val = q[0];
            int k = q[1];

            add(r, val);
            ++sz2;

            int v = r.firstKey();

            remove(r, v);
            --sz2;

            add(l, v);
            ++sz1;

            while (sz2 < k) {
                v = l.lastKey();

                remove(l, v);
                --sz1;

                add(r, v);
                ++sz2;
            }

            while (sz2 > k) {
                v = r.firstKey();

                remove(r, v);
                --sz2;

                add(l, v);
                ++sz1;
            }

            int x = r.firstKey();

            p = (int) qpow(p, x, mod);

            ans.add(p);
        }

        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
class Solution {
public:
    using ll = long long;

    void add(map<int, int>& mp, int x) {
        ++mp[x];
    }

    void remove(map<int, int>& mp, int x) {
        if (--mp[x] == 0) {
            mp.erase(x);
        }
    }

    ll qpow(ll a, int b, int mod) {
        ll ans = 1;

        while (b) {
            if (b & 1) {
                ans = ans * a % mod;
            }

            a = a * a % mod;
            b >>= 1;
        }

        return ans;
    }

    vector<int> powerUpdate(vector<int>& nums, int p, vector<vector<int>>& queries) {
        map<int, int> l, r;

        int sz1 = 0, sz2 = nums.size();

        for (int x : nums) {
            add(r, x);
        }

        const int mod = 1e9 + 7;

        vector<int> ans;

        for (auto& q : queries) {
            int val = q[0];
            int k = q[1];

            add(r, val);
            ++sz2;

            int v = r.begin()->first;

            remove(r, v);
            --sz2;

            add(l, v);
            ++sz1;

            while (sz2 < k) {
                v = l.rbegin()->first;

                remove(l, v);
                --sz1;

                add(r, v);
                ++sz2;
            }

            while (sz2 > k) {
                v = r.begin()->first;

                remove(r, v);
                --sz2;

                add(l, v);
                ++sz1;
            }

            int x = r.begin()->first;

            p = qpow(p, x, mod);

            ans.push_back(p);
        }

        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
func powerUpdate(nums []int, p int, queries [][]int) []int {
    l := redblacktree.New[int, int]()
    r := redblacktree.New[int, int]()

    merge := func(st *redblacktree.Tree[int, int], x, v int) {
        c, _ := st.Get(x)

        if c+v == 0 {
            st.Remove(x)
        } else {
            st.Put(x, c+v)
        }
    }

    sz1, sz2 := 0, len(nums)

    for _, x := range nums {
        merge(r, x, 1)
    }

    const mod int = 1e9 + 7

    qpow := func(a, b int) int {
        ans := 1

        for b > 0 {
            if b&1 == 1 {
                ans = ans * a % mod
            }

            a = a * a % mod
            b >>= 1
        }

        return ans
    }

    ans := make([]int, 0, len(queries))

    for _, q := range queries {
        val, k := q[0], q[1]

        merge(r, val, 1)
        sz2++

        node := r.Left()

        merge(r, node.Key, -1)
        sz2--

        merge(l, node.Key, 1)
        sz1++

        for sz2 < k {
            node = l.Right()

            merge(l, node.Key, -1)
            sz1--

            merge(r, node.Key, 1)
            sz2++
        }

        for sz2 > k {
            node = r.Left()

            merge(r, node.Key, -1)
            sz2--

            merge(l, node.Key, 1)
            sz1++
        }

        x := r.Left().Key

        p = qpow(p, x)

        ans = append(ans, p)
    }

    return ans
}

Solution 2: Sorted List

Thinking

Solution 1 keeps two ordered sets and shuttles by \(k\). The same limits also allow a single sorted list: after each insertion, \(sl[-k]\) is already the \(k\)-th largest.

The code is shorter, no longer relies on \(k\) changing slowly, and has a smaller implementation footprint.

We use a sorted list \(\textit{sl}\) to maintain the current array \(nums\). For each query, we insert \(val_i\) into \(\textit{sl}\), then find the \(k_i\)-th largest element \(x\) in \(\textit{sl}\). Using fast exponentiation, we update \(p\) to \(p^x \bmod (10^9 + 7)\), and append the updated \(p\) to the answer array.

The time complexity is \(O((n + m) \log (n + m))\), and the space complexity is \(O(n + m)\), where \(n\) and \(m\) are the lengths of \(\textit{nums}\) and \(\textit{queries}\), respectively.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution:
    def powerUpdate(
        self, nums: list[int], p: int, queries: list[list[int]]
    ) -> list[int]:
        ans = []
        sl = SortedList(nums)
        mod = 10**9 + 7
        for val, k in queries:
            sl.add(val)
            x = sl[-k]
            p = pow(p, x, mod)
            ans.append(p)
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <vector>

using namespace std;
using namespace __gnu_pbds;

template <typename T>
using ordered_multiset = tree<pair<T, int>, null_type, less<pair<T, int>>,
    rb_tree_tag, tree_order_statistics_node_update>;

class Solution {
public:
    vector<int> powerUpdate(vector<int>& nums, int p, vector<vector<int>>& queries) {
        vector<int> ans;
        ordered_multiset<int> sl;
        const int mod = 1e9 + 7;

        for (int i = 0; i < nums.size(); i++) {
            sl.insert({nums[i], i});
        }

        int next_id = nums.size();

        auto mod_pow = [&](long long base, long long exp) -> long long {
            long long result = 1;
            base %= mod;
            while (exp > 0) {
                if (exp & 1) result = (result * base) % mod;
                base = (base * base) % mod;
                exp >>= 1;
            }
            return result;
        };

        for (const auto& query : queries) {
            int val = query[0];
            int k = query[1];

            sl.insert({val, next_id++});

            auto it = sl.find_by_order(sl.size() - k);
            int kth_largest = it->first;

            p = mod_pow(p, kth_largest);
            ans.push_back(p);
        }

        return ans;
    }
};

Comments