本文正在参与掘金团队号上线活动，点击 查看大厂春招职位
题目描述
给你两个整数 m 和 k ，以及数据流形式的若干整数。你需要实现一个数据结构，计算这个数据流的 MK 平均值 。
MK 平均值 按照如下步骤计算：
如果数据流中的整数少于 m 个，MK 平均值 为 -1 ，否则将数据流中最后 m 个元素拷贝到一个独立的容器中。 从这个容器中删除最小的 k 个数和最大的 k 个数。 计算剩余元素的平均值，并 向下取整到最近的整数 。 请你实现 MKAverage 类：
MKAverage(int m, int k) 用一个空的数据流和两个整数 m 和 k 初始化 MKAverage 对象。 void addElement(int num) 往数据流中插入一个新的元素 num 。 int calculateMKAverage() 对当前的数据流计算并返回 MK 平均数 ，结果需 向下取整到最近的整数 。
提示：
3 <= m <= 10^5
1 <= k*2 < m
1 <= num <= 10^5
addElement 与 calculateMKAverage 总操作次数不超过 10^5 次。来源：力扣（LeetCode）
链接：leetcode-cn.com/problems/fi…
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
题目分析
如果直接计算，每次处理数据的时间复杂度为 O(M)，共有 10^5 次操作肯定会超时。上一篇题解说过，做题之前先看数据范围，此题 num 的数据不超过 10^5 可以考虑从此入手。
我们开一个数组，count[100000] 用来记录每一个数字出现的次数，这样我们取前 k 个数字和后 k 个数字都可以直接从前或从后开始把每个数字出现的次数加起来，直到为 k 即可。
简单用代码表示，取最小的 k 个数字的和。
let countSum = 0;
let valueSum = 0;
for (let i = 1; i <= 100000; i++) { // 数字的范围是 [1, 100000]
// i 出现 count[i] 次
valueSum += i * min(count[i], k - countSum); // 如果当前数字出现次数全部加起来大于k 注意不要加多
countSum += count[i];
if (countSum >= k) {
break;
}
}
return valueSum;
当然此时一次计算的复杂度还是 O(N) 仍然会超时。我们需要想办法降到 O(logN) ，此时可以想到一个数据结构——线段树。
线段树可以以 O(logN) 的时间复杂度求出区间和，进行点更新，区间更新。
线段树就是用空间换时间的典型应用，它是一颗二叉树。讲解线段树的文章很多，不了解的同学可以自己去搜索一下。我没把握讲清楚就不误人子弟了。:D
回到此题，线段树需要维护两个值，一个是区间出现的数字个数和，一个是区间出现数字的加和。
当增加一个数字 num 时，就是在线段树的 num 处，数字个数加一，数字加和加 num。
求平均值时，需要求前 k 个数的和 与 前 m-k 数的和，两个值求差就是中间段的数字和，然后再求平均值即可。
AC 代码
function SegTree(left, right) {
this.left = left; // 节点的左端点
this.right = right; // 节点的右端点
this.leftChild = null;
this.rightChild = null;
this.countSum = 0;
this.valueSum = 0;
}
SegTree.buildTree = function buildTree(left, right) {
const node = new SegTree(left, right);
if (left !== right) {
const mid = (left + right) >> 1;
node.leftChild = buildTree(left, mid);
node.rightChild = buildTree(mid + 1, right);
}
return node;
}
SegTree.prototype.isLeafNode = function() {
return !this.leftChild;
}
SegTree.prototype.includes = function(position) {
return this.left <= position && this.right >= position;
}
SegTree.prototype.pushUp = function() {
this.valueSum = this.leftChild.valueSum + this.rightChild.valueSum;
this.countSum = this.leftChild.countSum + this.rightChild.countSum;
}
// 在 position 处增加 addCount 个数字
SegTree.prototype.update = function(position, addCount) {
if (this.isLeafNode()) {
this.countSum += addCount;
this.valueSum += addCount * position;
return;
}
if (this.leftChild.includes(position)) {
this.leftChild.update(position, addCount);
} else {
this.rightChild.update(position, addCount);
}
this.pushUp();
}
// 查询前 count 个数的和
SegTree.prototype.query = function (count) {
if (this.isLeafNode()) return this.left * count;
if (this.countSum === count) return this.valueSum;
if (this.leftChild.countSum === count) return this.leftChild.valueSum;
if (this.leftChild.countSum < count)
return this.leftChild.valueSum + this.rightChild.query(count - this.leftChild.countSum);
return this.leftChild.query(count);
}
/**
* @param {number} m
* @param {number} k
*/
var MKAverage = function(m, k) {
// n <= 100000
this.m = m;
this.k = k;
this.segTree = SegTree.buildTree(1, 100000);
this.queue = [];
};
/**
* @param {number} num
* @return {void}
*/
MKAverage.prototype.addElement = function(num) {
this.queue.push(num);
this.segTree.update(num, 1);
if (this.queue.length > this.m) {
this.segTree.update(this.queue.shift(), -1);
}
};
/**
* @return {number}
*/
MKAverage.prototype.calculateMKAverage = function() {
const { m, k } = this;
if (this.queue.length < m) return -1;
const minK = this.segTree.query(k);
const minM_K = this.segTree.query(m - k);
return Math.floor((minM_K - minK) / (m - 2 * k));
};
其他
第一次尝试用 JavaScript 写线段树，虽然 AC 了 如果发现问题欢迎及时指出，谢谢。
