Leetcode 2897. Apply Operations on Array to Maximize Sum of Squares

• Leetcode 2897. Apply Operations on Array to Maximize Sum of Squares
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：2897. Apply Operations on Array to Maximize Sum of Squares

1. 解题思路

这一题事实上非常的简单，我们只需要想明白一些关键点就行了。

题中最终的目标是获得$k$个数，使得其平方和最大。因此，我们就只需要理解一下题中给出的操作会带来怎样的变化，以及要如何获取最大的平方和即可。

首先，对于题中给出的位操作，事实上只有在一种情况下才会产生变化，即$i=1,j=0$的情况下，会变为$i=0,j=1$，其余情况都不会产生变化。因此，我们无法改变各个位上$0,1$的数目总量，但是总可以通过一定的操作将其均匀化或者聚合到某几个数上。

然后，如果要使得两个数的平方和最大，我们考察某一位上的数最好是分布在较大的数还是在较小的数上，不妨令这个位上的数为$x$，然后其余位上的数分别为$a,b$，那么平方和就是：
$$(a+x{)}^2+b^2=a^2+2ax+x^2+b^2$$

显然，$a$较大时，可以获得更大的平方和。

综上，我们就可以获得我们最终的结论：

• 要取$k$个数，使得平方和最大，我们只需要统计每一位上拥有的$1$的个数，然后尽可能全部分配到这$k$个数当中即可。

2. 代码实现

给出最终的代码实现如下：

class Solution:
    def maxSum(self, nums: List[int], k: int) -> int:
        MOD = 10**9+7
        digits = [0 for _ in range(32)]
        flag = [pow(2, i, MOD) for i in range(32)]
        for x in nums:
            idx = 0
            while x != 0:
                digits[idx] += x % 2 
                x = x // 2
                idx += 1
        
        ans = 0
        for i in range(k):
            x = 0
            for idx in range(32):
                if digits[idx] > 0:
                    digits[idx] -= 1
                    x += flag[idx]
            ans = (ans + x * x) % MOD
        return ans

提交代码评测得到：耗时2510ms，占用内存30.6MB。