leetcode 周赛 506, 排名线段树

零、背景

这次比赛最后一题其实不难，我一开始就想到了解法，但是O(n^2 log(n) log(n)) 的复杂度被卡常数了。
结束比赛后，做了好几个优化，才通过。

本场题型概览如下。
A 题：循环。
B 题：统计+贪心。
C 题：贪心。
D 题：枚举+滑动窗口+二分+排名线段树。

一、判定好整数

题意：给一个数字，问各位数字平方之和是否比各位数字之和大于等于50。

思路：循环计算求和比大小即可

二、频率平衡子数组

题意：给一个数组，定义频率平衡子数组如下：
1）只包含一个元素，则是平衡的
2）至少两个元素，那么其中每个出现频率最高的元素，其出现次数都必须恰好是该子数组中其他每个不同值出现次数的两倍。
求最长频率平衡子数组的长度。

思路：滑动窗口+统计+贪心

根据定义2，可以发现，如果一个子数组是平衡的，元素值频率的频率集合大小必须为 2。

不过做这道题需要弄清题意。
因为这道题有问题，定义不完整。
例如对于出现多个元素，所有元素频率都相等的情况没有定义。

第二组样例是只有一个元素值，频率相同，答案是整个数组。
第三组样例是多个不同值元素，频率相同，答案却是1。

故，可以根据这两个样例反推出两个定义。
3）所有元素值相同，则是平衡的。
4）多个不同值频率相同，则是不平衡的。

定义梳理完整后，可以直接枚举所有子数组，借助前缀信息，边枚举边统计每个元素值的频率，以及频率的频率。

然后对统计数据分三种情况判断。
1）频率集合大小为1，即只有一个元素值，子数组是平衡的，更新答案。
2）频率的频率集合大小不为 2，不满足平衡的定义，跳过。
3）此时，频率的频率大小为2，得到最大值和最小值，判断是否满足2倍关系。

复杂度：O(n^2)

三、设备评分的最大和

题意：给一个n*m 的矩阵，每行代表一个设备，每列代表设备单元的容量。
现在可以对矩阵执行任意次操作：
1）选择之前未标记移除的行 i。
2）行 i 移除任意一个单元，并添加到任意行中。
3）标记行 i 为已移除过。
求最后所有行的最小单元容量的和（如果行没有单元，按0处理）。

思路：贪心

正规的方法来做这道题，会发现很难使用数据结构来表示。

观察题目特征，每一行求的是最小值，可以得到下面的分析和结论。

分析0：关注同一个元素
显然，同一个元素移动多次没有意义，等价于一步到位。
所以，可以确定贪心规则，每个元素最多移动一次。

分析1：关注移除的行 i。
移除某个单元时，如果移除的不是最小值，则操作后不影响这一行的最小值，但是有可能导致目标行的最小值更小。
由此得出结论，如果某一行有移除操作，肯定是这一行的最小值。

分析2：关注加入的目标行。
可以加入到任意行，如果加入后不是最小值，则不影响答案；
如果成为最小值，则会使答案更小。
显然，我们需要挑选某一行，加入后不会成为最小值。

加入哪一行才满足呢？
显然是加入第一列所有行中最小值所在的那一行，此时不影响答案。

分析3：万一最小值那一行也进行移除，最小值移动走了呢？
那就定义，最小值如果有移动，当做第一个操作。此后，就不会再移动了。

分析4：关注最大答案
分析2确定了，移动后，当前行答案为次大值，目标行答案不变。
那么最终答案也就确定了，除了最小值所在行取最小值，其他行都是次大值。

根据上述分析，可以得到相关贪心解法。

1）选择排序，得到最小值与次小值，复杂度O(n*m)。
2）找到最小值所在行
3）计算所有行次小值的和。
4）枚举最小值移动的行数，答案为最小值与其他行的次大值之和。

特殊情况：只有一列时，不操作答案最优。

四、至多 K 次交换后最大子数组和

题意：给一个数组，你可以选择任意两个元素进行交换，最多交换 K 次。
问最终可能得到的最大子数组和。
数据范围：n <= 1500

思路：枚举+滑动窗口+二分+排名线段树

这道题直接看，没有好的状态转移方程来表示交换行为。
所以只能枚举子数组，结合滑动窗口，来高效找到答案。

假设当前枚举的窗口为 [l,r)，则窗口外的区间为 [0,l) 和 [r,n)。

要使窗口内的元素和最大，显然是拿窗口内的最小值与窗口外的最大值进行交换，且需要满足窗口外的大于窗口内的，最多K次。
暴力做法复杂度：O(n^2*k)，显然会超时。

如果窗口内的元素与窗口外的元素都是有序的，则可以二分找到可以交换多少个元素。

此时问题转化为了：维护一个数据结构，可以动态添加与删除元素，然后查询 TOP K 的元素值。
这个是典型的平衡树问题，直接在带计数的平衡树上 Search 即可。

手动实现平衡树成本较高，比赛的时候，我是使用线段树上来求 TOP K 问题。

TOP K 要求有序，所以先对所有值进行排序与离散化（相同值通过下标区分）。
线段树中维护一个 count，代表一个值是否被添加进来。
显然，对于父节点，count 代表子树内值的个数。
由此，便可以在线段树中求 TOP K 问题。

// 返回有值的第 k 个元素与位置, rank 从 1 开始
pair<ll, int> Rank(int k, int l = 1, int r = maxNM, int rt = 1) {
  // assert(countVal[rt] >= k);
  if (l == r) {
    return vals[rt];
  }
  int m = (l + r) >> 1;
  if (countVal[rt << 1] >= k) {
    return Rank(k, lson);
  } else {
    return Rank(k - countVal[rt << 1], rson);
  }
}

回到二分这一步，比较条件是窗口内的前 mid 个元素是否大于窗口外的后 mid 个元素。

int l = 1, r = min(min(windowsSize, otherSize), K) + 1;  // [l, r)
while (l < r) {
  int mid = (l + r) >> 1;
  // window 的前 mid 个元素 都不大于 other 的后 mid 个元素
  if (segTreeWindows.Rank(mid).first <= segTreeOther.Rank(otherSize - mid + 1).first) {
    l = mid + 1;
  } else {
    r = mid;  // 不满足
  }
}
MyPrintf("l = %d, r = %d\n", l, r);
r--; // 左开右闭，需要减一

二分出实际交换的个数 k 后，还需要计算窗口内前 k 个元素和与窗口外的后 k 个元素和。
这里计算排名时，也返回下标，从而可以根据下标计算出区间和。

auto [val, idx] = segTreeWindows.Rank(r);
auto [val2, idx2] = segTreeOther.Rank(otherSize - r + 1);
ll sum = segTreeWindows.QuerySum(idx + 1, n) + segTreeOther.QuerySum(idx2, n);
ans = max(ans, sum);

注意事项：K 为 0 时，不进行交换，就是普通的最大子数组和问题。

剪枝：正数个数不大于 K 时，直接选择所有正整数。
剪枝：没有正整数时，直接选择最大的数字。

复杂度：O(n^2 log(n) log(n))

五、最后

这次比赛对我来说难度还好，但是目前没有模版，代码量比较大，加上卡常数，比赛期间没有通过四道题，排名73名。

最近也感冒了，分享 AI 相关的文章先暂停一周，这一周时间里分享一下线段树相关的算法。

《完》

-EOF-

本文公众号：天空的代码世界
个人微信号：tiankonguse
公众号 ID：tiankonguse-code