【有序集合 树状树状 二分查找】1825. 求出 MK 平均值

本文涉及知识点

有序集合 树状树状 二分查找 队列
二分查找算法合集

LeetCode1825. 求出 MK 平均值

给你两个整数 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 平均数 ，结果需 向下取整到最近的整数 。
示例 1：
输入：
[“MKAverage”, “addElement”, “addElement”, “calculateMKAverage”, “addElement”, “calculateMKAverage”, “addElement”, “addElement”, “addElement”, “calculateMKAverage”]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
输出：
[null, null, null, -1, null, 3, null, null, null, 5]
解释：
MKAverage obj = new MKAverage(3, 1);
obj.addElement(3); // 当前元素为 [3]
obj.addElement(1); // 当前元素为 [3,1]
obj.calculateMKAverage(); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
obj.addElement(10); // 当前元素为 [3,1,10]
obj.calculateMKAverage(); // 最后 3 个元素为 [3,1,10]
// 删除最小以及最大的 1 个元素后，容器为 [3]
// [3] 的平均值等于 3/1 = 3 ，故返回 3
obj.addElement(5); // 当前元素为 [3,1,10,5]
obj.addElement(5); // 当前元素为 [3,1,10,5,5]
obj.addElement(5); // 当前元素为 [3,1,10,5,5,5]
obj.calculateMKAverage(); // 最后 3 个元素为 [5,5,5]
// 删除最小以及最大的 1 个元素后，容器为 [5]
// [5] 的平均值等于 5/1 = 5 ，故返回 5
提示：
3 <= m <= 10⁵
1 <= k*2 < m
1 <= num <= 10⁵
addElement 与 calculateMKAverage 总操作次数不超过 10⁵ 次。

有序集合

有个有序集合，分别记录：最小k个元素，中间的m-k$\times$ 2个元素， 最大的k个元素。这个容易想到，但细节颇多。

树状数组 + 二分查找

int 型树状数组数组cnt记录，num的数量。
long long型树状数组sum记录其和。
增加元素：
num加到队列que中。
cnt.updata(num,1);
sum.updata(nun,num);
如果que.size() > m
cnt.updata(que.front(),-1)
sum.update(que.front(),-que.front());
que.pop()
求平均值：
如果cnt的总数量 小于m，直接返回-1。
二分查找第一个i1，使得cnt.sum(i1) >= k
最小k个数的和s1为：sum.sum(i1) - (cnt.sum(i1) -k)*i1
同理求得最小m-k个数的和s2。
返回：(s2-s1)/(m-2$\times$k)

代码

核心代码

template<class ELE = int >
class CTreeArr
{
public:
	CTreeArr(int iSize) :m_vData(iSize + 1)
	{

	}
	void Add(int index, ELE value)
	{
		index++;
		while (index < m_vData.size())
		{
			m_vData[index] += value;
			index += index & (-index);
		}
	}
	ELE Sum(int index)//[0...index]之和
	{
		index++;
		ELE ret = 0;
		while (index)
		{
			ret += m_vData[index];
			index -= index & (-index);
		}
		return ret;
	}
	ELE Sum() { return Sum(m_vData.size() - 2);	}
	ELE Get(int index)
	{
		return Sum(index) - Sum(index - 1);
	}
private:
	vector<ELE> m_vData;
};
template<class INDEX_TYPE>
class CBinarySearch
{
public:
	CBinarySearch(INDEX_TYPE iMinIndex, INDEX_TYPE iMaxIndex):m_iMin(iMinIndex),m_iMax(iMaxIndex) {}
	template<class _Pr>
	INDEX_TYPE FindFrist( _Pr pr)
	{
		auto left = m_iMin - 1;
		auto rightInclue = m_iMax;
		while (rightInclue - left > 1)
		{
			const auto mid = left + (rightInclue - left) / 2;
			if (pr(mid))
			{
				rightInclue = mid;
			}
			else
			{
				left = mid;
			}
		}
		return rightInclue;
	}
protected:
	const INDEX_TYPE m_iMin, m_iMax;
};
class MKAverage {
public:
	MKAverage(int m, int k) :m_cnt(100'000 + 1), m_sum(100'000 + 1) {
		this->m = m;
		this->k = k;
	}
	void addElement(int num) {
		m_que.emplace(num);
		m_cnt.Add(num, 1);
		m_sum.Add(num, num);
		if (m_que.size() > m) {
			m_cnt.Add(m_que.front(), -1);
			m_sum.Add(m_que.front(), -m_que.front());
			m_que.pop();
		}
	}
	int calculateMKAverage() {
		if (m_cnt.Sum() < m) { return -1; }
		long long i1 = MoreEqualCnt(k);
		long long i2 = MoreEqualCnt(m - k);
		long long s2 = m_sum.Sum(i2) - (m_cnt.Sum(i2) - (m - k)) * i2;
		long long s1 = m_sum.Sum(i1) - (m_cnt.Sum(i1) - k) * i1;
		return (s2 - s1) / (m - 2 * k);
	}
	int MoreEqualCnt(int cnt) {
		CBinarySearch bs(0, 100'000);
		return bs.FindFrist([&](int inx) {return m_cnt.Sum(inx) >= cnt; });
	}
	queue<int> m_que;
	CTreeArr<int> m_cnt;
	CTreeArr<long long> m_sum;
	int m, k;
};

单元测试

template<class T1, class T2>
void AssertEx(const T1& t1, const T2& t2)
{
	Assert::AreEqual(t1, t2);
}

template<class T>
void AssertEx(const vector<T>& v1, const vector<T>& v2)
{
	Assert::AreEqual(v1.size(), v2.size());
	for (int i = 0; i < v1.size(); i++)
	{
		Assert::AreEqual(v1[i], v2[i]);
	}
}

template<class T>
void AssertV2(vector<vector<T>> vv1, vector<vector<T>> vv2)
{
	sort(vv1.begin(), vv1.end());
	sort(vv2.begin(), vv2.end());
	Assert::AreEqual(vv1.size(), vv2.size());
	for (int i = 0; i < vv1.size(); i++)
	{
		AssertEx(vv1[i], vv2[i]);
	}
}

namespace UnitTest
{
	int k;
	vector<int> arrival,load;
	TEST_CLASS(UnitTest)
	{
	public:
		TEST_METHOD(TestMethod00)
		{
			MKAverage MK(3, 1);
			MK.addElement(3);
			MK.addElement(1);
			MK.addElement(10);
			MK.addElement(5);
			MK.addElement(5);
			MK.addElement(5);
		}
		TEST_METHOD(TestMethod01)
		{
			k = 3, arrival = { 1,2,3,4 }, load = { 1,2,1,2 };
			auto res = Solution().busiestServers(k, arrival, load);
			AssertEx({ 0 }, res);
		}
		TEST_METHOD(TestMethod02)
		{
			k = 3, arrival = { 1,2,3 }, load = { 10,12,11 };
			auto res = Solution().busiestServers(k, arrival, load);
			AssertEx({ 0 ,1,2}, res);
		}
		TEST_METHOD(TestMethod03)
		{
			k = 3, arrival = { 1,2,3,4,8,9,10 }, load = { 5,2,10,3,1,2,2 };
			auto res = Solution().busiestServers(k, arrival, load);
			AssertEx({1 }, res);
		}
		TEST_METHOD(TestMethod04)
		{
			k = 1, arrival = { 1 }, load = { 1 };
			auto res = Solution().busiestServers(k, arrival, load);
			AssertEx({ 0 }, res);
		}	
	};
}

扩展阅读

视频课程

先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。
https://edu.csdn.net/course/detail/38771

如何你想快速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程
https://edu.csdn.net/lecturer/6176

相关推荐

我想对大家说的话
《喜缺全书算法册》以原理、正确性证明、总结为主。
按类别查阅鄙人的算法文章，请点击《算法与数据汇总》。
有效学习：明确的目标 及时的反馈 拉伸区（难度合适） 专注
闻缺陷则喜(喜缺)是一个美好的愿望，早发现问题，早修改问题，给老板节约钱。
子墨子言之：事无终始，无务多业。也就是我们常说的专业的人做专业的事。
如果程序是一条龙，那算法就是他的是睛

测试环境

操作系统：win7 开发环境： VS2019 C++17
或者 操作系统：win10 开发环境： VS2022 C++17
如无特殊说明，本算法用**C++**实现。