跳转至

3935. 插入后第 K 大更新的幂 I 🔒

题目描述

给定一个整数数组 nums 和一个整数 p。

同时给定一个二维整数数组 queries，其中每个 queries[i] = [vali, ki] 并且 连续 的 ki 个值之间的差总是 小于 10。

对于每次查询：

  • 将 vali 插入到 nums。
  • 令 x 为当前 nums 中第 ki 个 最大 的元素。
  • 将 p 更新 为 px % (109 + 7)。

返回数组 ans，其中 ans[i] 表示在第 i 次查询后 p 的值。

 

示例 1：

输入：nums = [2], p = 4, queries = [[3,1],[1,2]]

输出：[64,4096]

解释：

i vali 当前
nums		 ki 第 ki
大		 p 新的 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

因此，ans = [64, 4096]。

示例 2：

输入：nums = [7,5], p = 6, queries = [[4,3],[7,2]]

输出：[1296,220296870]

解释：

i vali 当前
nums		 ki 第 ki
大		 p 新的 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

因此，ans = [1296, 220296870]。

 

提示：

  • 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
  • 对于 i > 0，有 |ki - ki - 1| < 10

解法

方法一：两个有序集合

我们用两个有序集合 \(l\) 和 \(r\) 来维护当前的数组 \(nums\)，其中 \(l\) 中的元素都不大于 \(r\) 中的元素，并且 \(r\) 中的元素个数等于 \(k_i\)。

对于每次查询，我们将 \(val_i\) 插入到 \(r\) 中，然后将 \(r\) 中最小的元素移动到 \(l\) 中，直到 \(r\) 中的元素个数等于 \(k_i\)。此时，\(r\) 中最小的元素就是当前 \(nums\) 中第 \(k_i\) 个最大的元素，我们利用快速幂算法将 \(p\) 更新为 \(p^x \bmod (10^9 + 7)\)，并将更新后的 \(p\) 加入答案数组中。

时间复杂度 \(O((n + m) \log (n + m))\)，空间复杂度 \(O(n + m)\)，其中 \(n\) 和 \(m\) 分别是 \(\textit{nums}\) 和 \(\textit{queries}\) 的长度。

 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
}

方法二：有序列表

我们用一个有序列表 \(\textit{sl}\) 来维护当前的数组 \(nums\)。对于每次查询，我们将 \(val_i\) 插入到 \(\textit{sl}\) 中，然后找到 \(\textit{sl}\) 中第 \(k_i\) 个最大的元素 \(x\)，利用快速幂算法将 \(p\) 更新为 \(p^x \bmod (10^9 + 7)\)，并将更新后的 \(p\) 加入答案数组中。

时间复杂度 \(O((n + m) \log (n + m))\)，空间复杂度 \(O(n + m)\)，其中 \(n\) 和 \(m\) 分别是 \(\textit{nums}\) 和 \(\textit{queries}\) 的长度。

 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;
    }
};

评论