2897. Apply Operations on Array to Maximize Sum of Squares
Description
You are given a 0-indexed integer array nums
and a positive integer k
.
You can do the following operation on the array any number of times:
- Choose any two distinct indices
i
andj
and simultaneously update the values ofnums[i]
to(nums[i] AND nums[j])
andnums[j]
to(nums[i] OR nums[j])
. Here,OR
denotes the bitwiseOR
operation, andAND
denotes the bitwiseAND
operation.
You have to choose k
elements from the final array and calculate the sum of their squares.
Return the maximum sum of squares you can achieve.
Since the answer can be very large, return it modulo 109 + 7
.
Example 1:
Input: nums = [2,6,5,8], k = 2 Output: 261 Explanation: We can do the following operations on the array: - Choose i = 0 and j = 3, then change nums[0] to (2 AND 8) = 0 and nums[3] to (2 OR 8) = 10. The resulting array is nums = [0,6,5,10]. - Choose i = 2 and j = 3, then change nums[2] to (5 AND 10) = 0 and nums[3] to (5 OR 10) = 15. The resulting array is nums = [0,6,0,15]. We can choose the elements 15 and 6 from the final array. The sum of squares is 152 + 62 = 261. It can be shown that this is the maximum value we can get.
Example 2:
Input: nums = [4,5,4,7], k = 3 Output: 90 Explanation: We do not need to apply any operations. We can choose the elements 7, 5, and 4 with a sum of squares: 72 + 52 + 42 = 90. It can be shown that this is the maximum value we can get.
Constraints:
1 <= k <= nums.length <= 105
1 <= nums[i] <= 109
Solutions
Solution 1: Bitwise Operation + Greedy
According to the problem description, for an operation, we can change $nums[i]$ to $nums[i] \textit{ AND } nums[j]$, and change $nums[j]$ to $nums[i] \textit{ OR } nums[j]$. Let's consider the bits of the numbers. If two bits are both $1$ or both $0$, the result of the operation will not change the bits. If two bits are different, the result of the operation will change the bits to $0$ and $1$, respectively. Therefore, we can move $1$ bits to $0$ bits, but not vice versa.
We can use an array $cnt$ to count the number of $1$ bits in each position, and then select $k$ numbers from them. To maximize the sum of squares, we should choose the largest numbers as much as possible. This is because, assuming the sum of squares of two numbers is $a^2 + b^2$ (where $a \gt b$), changing them to $(a + c)^2 + (b - c)^2 = a^2 + b^2 + 2c(a - b) + 2c^2 \gt a^2 + b^2$ will increase the sum of squares. Therefore, to maximize the sum of squares, we should choose the largest number.
The time complexity is $O(n \times \log M)$, and the space complexity is $O(\log M)$. Here, $M$ is the maximum value in the array.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
|
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 |
|
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 |
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
|