题目描述
小扣在秋日市集入口处发现了一个数字游戏。主办方共有 N
个计数器,计数器编号为 0 ~ N-1
。每个计数器上分别显示了一个数字,小扣按计数器编号升序将所显示的数字记于数组 nums
。每个计数器上有两个按钮,分别可以实现将显示数字加一或减一。小扣每一次操作可以选择一个计数器,按下加一或减一按钮。
主办方请小扣回答出一个长度为 N
的数组,第 i
个元素 (0 <= i < N)
表示将 0~i
号计数器 初始 所示数字操作成满足所有条件 nums[a]+1 == nums[a+1],(0 <= a < i)
的最小操作数。回答正确方可进入秋日市集。
由于答案可能很大,请将每个最小操作数对 1,000,000,007
取余。
样例
输入:nums = [3,4,5,1,6,7]
输出:[0,0,0,5,6,7]
解释:
i = 0,[3] 无需操作
i = 1,[3,4] 无需操作;
i = 2,[3,4,5] 无需操作;
i = 3,将 [3,4,5,1] 操作成 [3,4,5,6], 最少 5 次操作;
i = 4,将 [3,4,5,1,6] 操作成 [3,4,5,6,7], 最少 6 次操作;
i = 5,将 [3,4,5,1,6,7] 操作成 [3,4,5,6,7,8],最少 7 次操作;
返回 [0,0,0,5,6,7]。
输入:nums = [1,2,3,4,5]
输出:[0,0,0,0,0]
解释:对于任意计数器编号 i 都无需操作。
输入:nums = [1,1,1,2,3,4]
输出:[0,1,2,3,3,3]
解释:
i = 0,无需操作;
i = 1,将 [1,1] 操作成 [1,2] 或 [0,1] 最少 1 次操作;
i = 2,将 [1,1,1] 操作成 [1,2,3] 或 [0,1,2],最少 2 次操作;
i = 3,将 [1,1,1,2] 操作成 [1,2,3,4] 或 [0,1,2,3],最少 3 次操作;
i = 4,将 [1,1,1,2,3] 操作成 [-1,0,1,2,3],最少 3 次操作;
i = 5,将 [1,1,1,2,3,4] 操作成 [-1,0,1,2,3,4],最少 3 次操作;
返回 [0,1,2,3,3,3]。
限制
1 <= nums.length <= 10^5
1 <= nums[i] <= 10^3
算法
(动态中位数) $O(n \log n)$
- 将每个
nums[i]
都减去i
,题目可以转化为求将前缀都变为相等数字时的最小代价。 - 根据数轴法,给定若干个数字,将其都变为相等时的目标值一定是整个数列的中位数(当没有中位数时,中间两个数字或之间的任意值都行)。
- 所以题目变为了求所有前缀的中位数,即动态中位数。可以使用两个堆维护求解,除此之外还需要维护每个堆中所有数字的总和。
时间复杂度
- 每个数字插入堆中的时间复杂度为 $O(\log n)$。
- 故总时间复杂度为 $O(n \log n)$。
空间复杂度
- 需要 $O(n)$ 的额外空间存储两个堆。
C++ 代码
#define LL long long
class Solution {
public:
vector<int> numsGame(vector<int>& nums) {
const int mod = 1000000007;
const int n = nums.size();
priority_queue<int> ma;
priority_queue<int, vector<int>, greater<int>> mi;
LL sum_ma = 0, sum_mi = 0;
vector<int> ans(n);
ans[0] = 0;
ma.push(nums[0]);
sum_ma += nums[0];
for (int i = 1; i < n; i++) {
if (nums[i] - i <= ma.top()) {
ma.push(nums[i] - i);
sum_ma += nums[i] - i;
} else {
mi.push(nums[i] - i);
sum_mi += nums[i] - i;
}
if (ma.size() > mi.size() + 1) {
int x = ma.top();
ma.pop(); mi.push(x);
sum_ma -= x; sum_mi += x;
} else if (ma.size() < mi.size()) {
int x = mi.top();
mi.pop(); ma.push(x);
sum_mi -= x; sum_ma += x;
}
LL m = ma.top();
ans[i] = ((m * ma.size() - sum_ma) +
(sum_mi - m * mi.size())) % mod;
}
return ans;
}
};