LeetCode 1825 - 求出 MK 平均值

发表于 更新于 分类于 阅读次数：

题目链接：LeetCode 1825 - 求出 MK 平均值（第 236 场周赛第四题）

思路

我们需要动态维护一个长度为 \(m\) 的容器，当容器的长度超过 \(m\) 时，根据先进先出的原则对元素进行淘汰。这一点很自然想到使用队列。

我们要对容器中的元素统计第 \(k + 1\) 个到第 \(m - k\) 个的和，这显然是一个区间询问，我们需要选择一种数据结构，于是看怎么修改。这里之后添加和淘汰需要用到修改操作，是单点修改（加 / 减），于是我们可以用树状数组来解决这个问题。

我们定义两个树状数组：\(\rm cnt\) 和 \(\rm sum\)，\(\rm cnt\) 用来维护元素个数的前缀和，\(\rm sum\) 用来维护元素的前缀和。当我们需要获得第 \(k\) 个元素的时候，可以现在 \(\rm cnt\) 中二分第 \(k\) 个元素是什么，然后用 \(\rm sum\) 来统计答案。

代码

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
90
91
92
93
94
95
96
97
using LL = long long;

const int MAX_N = int(1E5) + 5;

struct BinaryIndexedTree {
    LL c[MAX_N];
    int maxPos;
    
    BinaryIndexedTree() {}
    BinaryIndexedTree(int maxPos) {
        this->maxPos = maxPos;
        for (int i = 0; i <= maxPos; ++i) c[i] = 0;
    }

    inline int lowbit(const int &u) {
        return u & (-u);
    }

    void add(int pos, int k) {
        while (pos <= maxPos) {
            c[pos] += k; pos += lowbit(pos);
        }
    }

    LL sum(int pos) {
        LL ret = 0;
        while (pos >= 1) {
            ret += c[pos]; pos -= lowbit(pos);
        }
        return ret;
    }
};

class MKAverage {
public:
    BinaryIndexedTree cnt, sum;
    int node_tot, m, k;
    queue<int> q;
    
    MKAverage(int m_, int k_) {
        cnt = BinaryIndexedTree(1E5);
        sum = BinaryIndexedTree(1E5);
        node_tot = 0;
        m = m_;
        k = k_;
    }
    
    void addElement(int num) {
        cnt.add(num, 1);
        sum.add(num, num);
        q.push(num);
        ++node_tot;

        if (node_tot > m) {
            auto f = q.front();
            q.pop();
            cnt.add(f, -1);
            sum.add(f, -f);
            --node_tot;
        }
    }
    
    // 第 p 个元素在 BIT 中的 index
    int pre_index(int p) {
        int l = 1, r = 1E5;
        while (l < r) {
            int mid = (l + r) >> 1;
            if (cnt.sum(mid) < p) {
                l = mid + 1;
            } else {
                r = mid;
            }
        }
        
        return l;
    } 
    
    int calculateMKAverage() {
        if (node_tot < m) {
            return -1;
        }
        
        int l = pre_index(k);
        int r = pre_index(node_tot - k);
        
        if (l == r) {
            return l;
        }
        
        LL ans = sum.sum(r - 1) - sum.sum(l);
        ans += 1LL * l * (cnt.sum(l) - k);
        ans += 1LL * r * (node_tot - k - cnt.sum(r - 1));
        ans = (ans / (node_tot - 2 * k));
        
        return ans;
    }
};

复杂度

记询问次数为 \(q\)，\(\rm num\) 和 \(m\) 的最大值都是 \(n\)（即本题的 \(10^5\)）。

  • 时间复杂度：\(O(q \times (\log n)^2)\)。\((\log n)^2\) 表示在树状数组上二分的复杂度。
  • 空间复杂度：\(O(q + n)\)。