从 LeetCode 2897 看「可反复操作」类问题：找不变量，直接构造终态

发表于 分类于 Waline：
从每一位上 1 的个数入手，解释为什么要集中分配二进制位，以及如何跳过模拟直接构造最优解。

最近的每日一题是 LeetCode 2897「对数组执行操作使平方和最大」。这道题的代码不长，但很适合练习处理「操作可以执行任意多次」的问题：先找不变量，再想清最优终态，最后直接构造。

题目

LeetCode 2897：对数组执行操作使平方和最大（原题）

给定数组 nums 和正整数 k。每次选择两个不同的下标 i、j，同时把 nums[i] 更新为旧值的 nums[i] & nums[j]，把 nums[j] 更新为旧值的 nums[i] | nums[j]。操作次数不限。最后选出 k 个元素，使平方和最大，并将结果对 10^9 + 7 取模。

约束是 1 <= k <= nums.length <= 10^5，1 <= nums[i] <= 10^9。

例如 nums = [2, 6, 5, 8]、k = 2 时，可以得到 15 和 6，答案是 15² + 6² = 261。

第一步：找操作下的不变量

看到「可以操作任意多次」，先不要模拟，先问：操作前后什么不会变？

把 AND 和 OR 按二进制位拆开：

a 的某一位b 的同一位ANDOR
0000
1111
1001
0101

无论是哪一行，操作前后这一位上 1 的总数都相同。操作只会改变 1 在哪些元素里，每一位上 1 的个数始终不变。

于是，问题可以先放宽成：如果每一位上 1 的个数固定，应该怎样把这些 1 分配给数组元素，才能让选出的 k 个数平方和最大？

第二步：确定贪心方向

在总和固定时，数值越集中，平方和越大。例如 4² + 0² = 16，而 2² + 2² = 8。

具体到某个值为 v = 2^t 的二进制位：设 a >= b，a 没有这一位，b 有这一位。把这个 1 从 b 分给 a 后，平方和增加：

(a+v)2+(b−v)2−a2−b2=2v(a−b)+2v2>0(a+v)^2+(b-v)^2-a^2-b^2=2v(a-b)+2v^2>0

如果 b 没有被选中，分给已选中的 a 更是只会增加目标值。因此，对于「只固定每位 1 的个数」这个放宽后的问题，最优分配会尽量集中：第一个数拿到所有出现过的位，第二个数拿到各位剩下的一份，以此类推。第 t 个构造出的数，恰好包含所有计数至少为 t 的位。

不过，这里还差一个关键问题：这样的分配真的能用题目允许的 AND/OR 操作得到吗？

可以。固定一个下标作为第一个数，让它依次和其余所有数做 OR；每次另一个数接收 AND。这样第一个数最终就是所有原数组元素的按位 OR，也就是每个出现过的位都拿到一个 1。由于每次操作都保持各位计数，剩余元素中每位的计数正好减去一。接着只在剩余元素里重复这个过程，就能依次造出上述第二个数、第三个数……所以这个最优分配是可达的，不只是一个理论上界。

以 [2, 6, 5, 8] 为例，从低位到高位，四个位的计数分别是 [1, 2, 2, 1]。第一轮取到 1111₂ = 15，第二轮取到 0110₂ = 6，前两个数的平方和就是 261。

第三步：不模拟，直接构造

既然已经知道终态，只需统计每一位有多少个 1，再造出前 k 个数。每造一个数，就从仍有剩余计数的每一位各取一个 1。

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
var maxSum = function (nums, k) {
  const MOD = 1_000_000_007n;
  // nums[i] <= 10^9 < 2^30，只需处理第 0～29 位。
  const cnt = new Array(30).fill(0);

  for (const num of nums) {
    for (let bit = 0; bit < 30; bit++) {
      if ((num >> bit) & 1) cnt[bit]++;
    }
  }

  let ans = 0n;
  for (let i = 0; i < k; i++) {
    let x = 0;
    for (let bit = 0; bit < 30; bit++) {
      if (cnt[bit] > 0) {
        x |= 1 << bit;
        cnt[bit]--;
      }
    }
    const value = BigInt(x);
    ans = (ans + value * value) % MOD;
  }

  return Number(ans);
};

这里有两个实现细节：

  • 一个构造出的数最大接近 2^30，平方可能接近 10^18，超出 JavaScript 的安全整数范围，所以平方与取模要使用 BigInt。最终余数小于 10^9 + 7，转回 Number 是安全的。
  • 只需要各位的计数，不依赖原数组顺序，因此不需要排序，也不需要真的执行 AND/OR 操作。

时间复杂度为 O(30(n + k))，额外空间复杂度为 O(30)，其中 n = nums.length。

从这道题带走什么

  1. 先找不变量：能反复执行操作时，先观察操作保留了哪些量。
  2. 按位拆开：遇到 AND、OR 等位运算，先逐位分析，再看各位能否独立计数。
  3. 根据目标函数判断方向：总和固定时，平方和最大倾向于集中，平方和最小倾向于压平。
  4. 证明终态可达后直接构造：只找到了不变量，还不足以断言任意分配都能由合法操作实现。

前几天的 LeetCode 2333「最小差值平方和」 正好提供了相反方向的例子：它要最小化平方和，因此优先削减较大的差值；这题则要把各位的 1 尽量集中。