对数组执行操作使平方和最大

标签: 贪心 位运算 数组 哈希表

难度: Hard

给你一个下标从 0 开始的整数数组 nums 和一个 正 整数 k 。

你可以对数组执行以下操作 任意次 ：

  • 选择两个互不相同的下标 i 和 j ，同时 将 nums[i] 更新为 (nums[i] AND nums[j]) 且将 nums[j] 更新为 (nums[i] OR nums[j]) ，OR 表示按位 或 运算，AND 表示按位 与 运算。

你需要从最终的数组里选择 k 个元素，并计算它们的 平方 之和。

请你返回你可以得到的 最大 平方和。

由于答案可能会很大，将答案对 109 + 7 取余 后返回。

示例 1：

输入：nums = [2,6,5,8], k = 2
输出：261
解释：我们可以对数组执行以下操作：
- 选择 i = 0 和 j = 3 ，同时将 nums[0] 变为 (2 AND 8) = 0 且 nums[3] 变为 (2 OR 8) = 10 ，结果数组为 nums = [0,6,5,10] 。
- 选择 i = 2 和 j = 3 ，同时将 nums[2] 变为 (5 AND 10) = 0 且 nums[3] 变为 (5 OR 10) = 15 ，结果数组为 nums = [0,6,0,15] 。
从最终数组里选择元素 15 和 6 ，平方和为 152 + 62 = 261 。
261 是可以得到的最大结果。

示例 2：

输入：nums = [4,5,4,7], k = 3
输出：90
解释：不需要执行任何操作。
选择元素 7 ，5 和 4 ，平方和为 72 + 52 + 42 = 90 。
90 是可以得到的最大结果。

提示：

  • 1 <= k <= nums.length <= 105
  • 1 <= nums[i] <= 109

Submission

运行时间: 1189 ms

内存: 29.6 MB

"""
nums[i] = 0, nums[j] = 0，无变化
nums[i] = 0, nums[j] = 1，可以交换 0 和 1 的位置，也可以不变
nums[i] = 1, nums[j] = 1，无变化

所以对数组的操作就是交互两个数某个比特位上的0和1，但两个数的和是不变的
(a+b)^2 = a^2 + b^2 + 2ab
a^2+b^2 = (a+b)^2 - 2ab，a+b的值不变，要使平方和最大，需要使 ab 最小，a 和 b 的差要最大，所以要把 1 集中到一个数上
"""

class Solution:
    def maxSum(self, nums: List[int], k: int) -> int:
        m = max(nums).bit_length()
        f = [0] * m
        for i in range(m):
            f[i] = sum(x >> i & 1 for x in nums)
        ans = 0
        mod = 10 ** 9 + 7
        while True:
            x = 0
            for i in range(m):
                if f[i]:
                    x += (1 << i)
                    f[i] -= 1
            ans = (ans + x * x) % mod
            k -= 1
            if x == 0 or k == 0:
                break
        return ans

Explain

该题解首先获取数组中最大数的二进制位长度。然后，创建一个数组 f，其长度为最大二进制位长度，用于统计 nums 中每个比特位上 1 的个数。随后，采用贪心算法尝试构造最大的数。具体来说，从最高位到最低位，如果某个位上有 1，就将该位加到构造的数 x 中，并更新计数。这样构建的 x 将会是局部最大的可能值。重复这一过程 k 次或直到无法构造更多的非零数 x。每次迭代中计算 x 的平方，并累加到结果 ans 中。最终，返回的 ans 就是 k 个数的平方和最大的可能值。

时间复杂度: O((n+k)*m)

空间复杂度: O(m)

"""
nums[i] = 0, nums[j] = 0，无变化
nums[i] = 0, nums[j] = 1，可以交换 0 和 1 的位置，也可以不变
nums[i] = 1, nums[j] = 1，无变化

所以对数组的操作就是交互两个数某个比特位上的0和1，但两个数的和是不变的
(a+b)^2 = a^2 + b^2 + 2ab
a^2+b^2 = (a+b)^2 - 2ab，a+b的值不变，要使平方和最大，需要使 ab 最小，a 和 b 的差要最大，所以要把 1 集中到一个数上
"""

class Solution:
    def maxSum(self, nums: List[int], k: int) -> int:
        m = max(nums).bit_length()  # 获取 nums 中最大值的位数
        f = [0] * m  # f[i] 将存储每个位上 1 的总个数
        for i in range(m):  # 计算每个位上 1 的个数
            f[i] = sum(x >> i & 1 for x in nums)
        ans = 0  # 初始化结果变量
        mod = 10 ** 9 + 7  # 定义模数
        while True:
            x = 0  # 从 f 中构建最大数
            for i in range(m):  # 对于每个位，尝试构建最大数
                if f[i]:
                    x += (1 << i)  # 构建数
                    f[i] -= 1  # 减少该位的可用 1 的数量
            ans = (ans + x * x) % mod  # 更新结果
            k -= 1  # 减少还需要构建的数的数量
            if x == 0 or k == 0:  # 如果 x 为 0 或已达到 k 个数，结束循环
                break
        return ans  # 返回最大平方和

Explore

题解中描述的操作实际上是不会形成 k 个最大数值的确切解决方案。这种操作仅仅是对位操作的一种思考，其目的在于通过组合 nums[i] 和 nums[j] 的位来尝试形成更大的数值。然而，这种操作并不直接关注于形成最大的 k 个数值，而是更多的是通过位操作的方式尝试增加某些数的大小。真正的解决方案是使用贪心算法从数组 f 中构建最大可能的数，这种方法更直接有效地针对问题的目标。

在构造最大数值的过程中，从最高位到最低位开始是因为高位的数值权重远大于低位。例如，在二进制表示中，最高位（如第 31 位）的权重是 2^31，远大于第 1 位的权重 2^0。因此，为了使得构造的数值尽可能大，应优先考虑在高位放置尽可能多的 1。虽然低位的 1 的数量可能较多，但它们对最终数值的影响较小，因此从最高位开始构造是最有效的策略。

数组 f 的长度由 nums 中最大数的二进制位数决定，通常为 32 或 64，这取决于数据类型（如 int 或 long）。因此，f 的内存占用是非常小的，即使对于非常大的 nums 数组，f 数组的大小也不会变得很大。此外，对于每个位的计算，虽然需要遍历整个 nums 数组，但这是线性时间复杂度的操作，通常对于现代计算系统来说是可接受的。总的来说，这种方法在空间和时间上都是高效的，适用于处理大规模数据。