Problem

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 and j and simultaneously update the values of nums[i] to (nums[i] AND nums[j]) and nums[j] to (nums[i] OR nums[j]). Here, OR denotes the bitwise OR operation, and AND denotes the bitwise AND operation.

You have to choose k elements from the final array and calculate the sum of their squares.

Return themaximum sum of squares you can achieve.

Since the answer can be very large, return it modulo 10^9 + 7.

Examples

Example 1

1
2
3
4
5
6
7

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

1
2
3
4
5

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 <= 10^5
  • 1 <= nums[i] <= 10^9

Solution

Method 1 – Bitwise Greedy Construction

Intuition

The key insight is that the allowed operation lets us “move” all set bits in the array into a single element, maximizing its value. By greedily collecting the most significant bits into up to k elements, we can maximize the sum of their squares.

Approach

  1. Count the number of times each bit (from 0 to 30) is set across all numbers in nums.
  2. For each of the k elements we want to maximize, construct a number by greedily assigning the highest available bits (using up the bit counts).
  3. For each constructed number, square it and sum up the results.
  4. Return the sum modulo 10^9 + 7.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class Solution {
public:
  int maxSumSquares(vector<int>& nums, int k) {
    const int MOD = 1e9 + 7;
    vector<int> bitCnt(31, 0);
    for (int x : nums)
      for (int b = 0; b < 31; ++b)
        if (x & (1 << b)) bitCnt[b]++;
    vector<long long> vals(k, 0);
    for (int b = 30; b >= 0; --b) {
      int cnt = bitCnt[b];
      for (int i = 0; i < k && cnt > 0; ++i, --cnt)
        vals[i] |= (1LL << b);
    }
    long long ans = 0;
    for (long long v : vals)
      ans = (ans + v * v % MOD) % MOD;
    return (int)ans;
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

class Solution {
  public int maxSumSquares(int[] nums, int k) {
    int MOD = (int)1e9 + 7;
    int[] bitCnt = new int[31];
    for (int x : nums)
      for (int b = 0; b < 31; ++b)
        if ((x & (1 << b)) != 0) bitCnt[b]++;
    long[] vals = new long[k];
    for (int b = 30; b >= 0; --b) {
      int cnt = bitCnt[b];
      for (int i = 0; i < k && cnt > 0; ++i, --cnt)
        vals[i] |= (1L << b);
    }
    long ans = 0;
    for (long v : vals)
      ans = (ans + v * v % MOD) % MOD;
    return (int)ans;
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class Solution:
  def maxSumSquares(self, nums: list[int], k: int) -> int:
    MOD = 10**9 + 7
    bit_cnt: list[int] = [0] * 31
    for x in nums:
      for b in range(31):
        if x & (1 << b):
          bit_cnt[b] += 1
    vals: list[int] = [0] * k
    for b in range(30, -1, -1):
      cnt = bit_cnt[b]
      for i in range(k):
        if cnt == 0:
          break
        vals[i] |= (1 << b)
        cnt -= 1
    ans: int = 0
    for v in vals:
      ans = (ans + v * v) % MOD
    return ans

Complexity

  • ⏰ Time complexity: O(31 * n + 31 * k) (since 31 bits per number and up to k elements)
  • 🧺 Space complexity: O(31 + k)