# Apply Operations on Array to Maximize Sum of Squares
**Difficulty:** HARD
[External](https://leetcode.com/problems/apply-operations-on-array-to-maximize-sum-of-squares)
Canonical: https://scaleengineer.com/dsa/problems/apply-operations-on-array-to-maximize-sum-of-squares
**Patterns:** [Greedy](https://scaleengineer.com/dsa/patterns/greedy), [Bit Manipulation](https://scaleengineer.com/dsa/patterns/bit-manipulation)
**Data structures:** Array, Hash Table
**Companies:** [Sprinklr](https://scaleengineer.com/companies/sprinklr)
---
## 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 _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`

# Approaches
## Greedy Bit Distribution with Sorting
This approach is based on the crucial observation that the bitwise operations `(a, b) -> (a&b, a|b)` conserve the total count of set bits at each position across the array. This insight allows us to treat all bits as being in a collective pool. The strategy is to count all set bits at each position, then use this pool of bits to reconstruct the numbers in the array to maximize the sum of squares of `k` elements. This is achieved by making `k` numbers as large as possible. This approach reconstructs the first `k` numbers greedily, then sorts the entire array to identify the top `k` elements for the final calculation.
**Time:** O(N*M + N log N), where N is the number of elements in `nums` and M is the number of bits (~31). The bit counting takes O(N*M). Reconstructing the numbers takes O(k*M). Sorting takes O(N log N). The dominant term is `N log N` or `N*M` depending on their relative sizes, but typically sorting is the bottleneck. · **Space:** O(N), where N is the number of elements in `nums`. This is for storing the `newNums` list. The `bitCounts` array takes O(M) space, where M is the number of bits (a constant, ~31), so it's dominated by O(N).
**Pros:** The approach is guaranteed to find the correct maximum sum.; The logic of pooling and redistributing bits is sound and correctly captures the essence of the problem.
**Cons:** The sorting step `O(N log N)` is computationally unnecessary and makes this approach less efficient than the optimal one.; It requires creating and manipulating an auxiliary array of size `N`, which might use more memory than necessary if `N` is very large.
### Explanation
The fundamental idea is that we can rearrange the bits among the numbers in the array without losing any. To maximize the sum of squares (`Σxᵢ²`), we should make the numbers as unequal as possible (e.g., `10² + 0² > 5² + 5²`). This means we should concentrate the most valuable (most significant) bits into a few numbers.

This approach implements this by:
1. Counting the total number of set bits for each bit position (0 to 30) over all numbers in `nums`.
2. Creating a new array of size `n`.
3. Greedily constructing `k` large numbers in this new array using the counted bits. We iterate from the most significant bit downwards, distributing the available bits to the first `k` slots.
4. Sorting this new array to find the `k` largest numbers.
5. Summing the squares of these `k` numbers modulo `10^9 + 7`.

While correct, the sorting step is a bottleneck that can be optimized away.

```java
import java.util.Collections;
import java.util.ArrayList;
import java.util.List;

class Solution {
    public int maxSum(List<Integer> nums, int k) {
        int n = nums.size();
        int MOD = 1_000_000_007;
        int[] bitCounts = new int[31];

        // 1. Count bits for each position
        for (int num : nums) {
            for (int i = 0; i < 31; i++) {
                if (((num >> i) & 1) == 1) {
                    bitCounts[i]++;
                }
            }
        }

        // 2. Reconstruct the first k numbers to be as large as possible
        List<Long> newNums = new ArrayList<>();
        for(int i = 0; i < n; i++) newNums.add(0L);

        for (int i = 30; i >= 0; i--) {
            for (int j = 0; j < k; j++) {
                if (bitCounts[i] > 0) {
                    newNums.set(j, newNums.get(j) + (1L << i));
                    bitCounts[i]--;
                } else {
                    break;
                }
            }
        }

        // 3. Sort to find the k largest elements
        Collections.sort(newNums, Collections.reverseOrder());

        // 4. Calculate sum of squares of the top k elements
        long totalSumOfSquares = 0;
        for (int i = 0; i < k; i++) {
            long num = newNums.get(i);
            long term = num % MOD;
            totalSumOfSquares = (totalSumOfSquares + (term * term)) % MOD;
        }

        return (int) totalSumOfSquares;
    }
}
```
### Algorithm
1. **Analyze the Operation**: The core of the problem lies in understanding the bitwise operation `(nums[i], nums[j]) -> (nums[i] & nums[j], nums[i] | nums[j])`. A key property of this operation is that it preserves the sum of the two numbers involved (`a + b = (a & b) + (a | b)`). More importantly, it preserves the total count of set bits at each bit position across the entire array. This means we can treat all the set bits from all numbers as being in a single pool, ready to be redistributed.
2. **Count Bits**: Create an integer array, `bitCounts`, of size 31 (since `nums[i] <= 10^9 < 2^30`). Iterate through each number in the input `nums` list. For each number, iterate through its bits from 0 to 30. If a bit is set, increment the corresponding counter in `bitCounts`.
3. **Reconstruct and Sort**: After counting, create a new array, `newNums`, of the same size as the input. The goal is to create `k` very large numbers. Distribute the collected bits into the first `k` elements of `newNums`. Iterate from the most significant bit (30) down to 0. For each bit position, distribute the available bits (given by `bitCounts`) one by one to `newNums[0]`, `newNums[1]`, ..., up to `newNums[k-1]`. The remaining `n-k` elements of `newNums` will be zero.
4. **Find k Largest**: Although the first `k` elements are constructed to be large, they are not necessarily sorted. Sort the `newNums` array in descending order to easily pick the `k` largest values.
5. **Calculate Sum of Squares**: Take the first `k` elements from the sorted `newNums` array. Calculate the sum of their squares, taking care to apply the modulo operation at each step to prevent overflow.

## Optimal Greedy Construction from Bit Counts
This approach refines the previous one by eliminating the unnecessary sorting step. It leverages the same core insight about the conservation of bit counts. After counting the bits, it directly constructs the `k` largest possible numbers without needing to manipulate a full-sized array. The greedy strategy for distributing bits ensures that the constructed numbers are the exact `k` numbers that yield the maximum sum of squares, thus making a final sorting step redundant. This leads to a more efficient solution.
**Time:** O((N+k)*M), where N is `nums.length`, M is the number of bits (~31). Counting bits takes O(N*M). Constructing the `k` numbers takes O(k*M). Since M is a small constant, the complexity is effectively linear in N and k. · **Space:** O(M + k), where M is the number of bits (~31) and k is the number of elements to choose. O(M) for `bitCounts` and O(k) for `resultNums`.
**Pros:** This is the most efficient solution with a linear time complexity.; It avoids unnecessary operations like sorting and uses space proportional to `k` rather than `N` for the constructed numbers.
**Cons:** The logic might be slightly less intuitive at first glance compared to an explicit sort.
### Explanation
The key to optimality is realizing that the sorting step is not needed. The process of constructing the numbers can be designed to produce the `k` largest values directly. By distributing the most significant bits first, we ensure the numbers are built in descending order of magnitude.

The algorithm is as follows:
1. Count the total set bits for each position (0-30) into a `bitCounts` array.
2. Create a `long` array `resultNums` of size `k`.
3. Iterate from bit position `i = 30` down to `0`. For each position, distribute the `bitCounts[i]` available bits among the `k` numbers in `resultNums`. The first bit goes to `resultNums[0]`, the second to `resultNums[1]`, and so on, in a round-robin fashion for that bit position. This ensures `resultNums[0] >= resultNums[1] >= ...`.
4. Finally, compute the sum of squares of the numbers in `resultNums` modulo `10^9 + 7`.

This avoids the `O(N log N)` sorting cost, resulting in a faster linear time solution.

```java
import java.util.List;

class Solution {
    public int maxSum(List<Integer> nums, int k) {
        int MOD = 1_000_000_007;
        int[] bitCounts = new int[31];

        // 1. Count bits for each position
        for (int num : nums) {
            for (int i = 0; i < 31; i++) {
                if (((num >> i) & 1) == 1) {
                    bitCounts[i]++;
                }
            }
        }

        // 2. Directly construct the k largest numbers
        long[] resultNums = new long[k];
        for (int i = 30; i >= 0; i--) { // Iterate from MSB to LSB
            for (int j = 0; j < k; j++) {
                if (bitCounts[i] > 0) {
                    resultNums[j] += (1L << i);
                    bitCounts[i]--;
                } else {
                    // No more bits at this position, move to the next bit
                    break;
                }
            }
        }

        // 3. Calculate sum of squares
        long totalSumOfSquares = 0;
        for (long num : resultNums) {
            long term = num % MOD;
            totalSumOfSquares = (totalSumOfSquares + (term * term)) % MOD;
        }

        return (int) totalSumOfSquares;
    }
}
```
### Algorithm
1. **Count Bits**: Same as the previous approach, count the total number of set bits for each position from 0 to 30 across all numbers in `nums`. Store these in a `bitCounts` array of size 31.
2. **Directly Construct `k` Largest Numbers**: Instead of modifying the original array or creating a new one of size `n`, create a result array `resultNums` of size `k`. Initialize all its elements to zero.
3. **Greedy Bit Distribution**: Iterate from the most significant bit (position 30) down to 0. For each bit position `i`, you have `bitCounts[i]` bits of value `2^i` to distribute. Distribute these bits one by one to `resultNums[0]`, `resultNums[1]`, ..., up to `resultNums[k-1]`. This ensures that `resultNums[0]` becomes the largest number, `resultNums[1]` the second largest, and so on, because they receive the highest-value bits first.
4. **Calculate Sum of Squares**: After distributing all bits, the `resultNums` array contains the `k` numbers that maximize the sum of squares. Iterate through this array, square each number, and add it to a running total. Remember to apply the modulo `10^9 + 7` at each addition to prevent overflow.

# Solutions
### Java

```java
class Solution {
public
  int maxSum(List<Integer> nums, int k) {
    final int mod = (int)1 e9 + 7;
    int[] cnt = new int[31];
    for (int x : nums) {
      for (int i = 0; i < 31; ++i) {
        if ((x >> i & 1) == 1) {
          ++cnt[i];
        }
      }
    }
    long ans = 0;
    while (k-- > 0) {
      int x = 0;
      for (int i = 0; i < 31; ++i) {
        if (cnt[i] > 0) {
          x |= 1 << i;
          --cnt[i];
        }
      }
      ans = (ans + 1L * x * x) % mod;
    }
    return (int)ans;
  }
}

```

### CPP

```cpp
class Solution {
public:
  int maxSum(vector<int> &nums, int k) {
    int cnt[31]{};
    for (int x : nums) {
      for (int i = 0; i < 31; ++i) {
        if (x >> i & 1) {
          ++cnt[i];
        }
      }
    }
    long long ans = 0;
    const int mod = 1e9 + 7;
    while (k--) {
      int x = 0;
      for (int i = 0; i < 31; ++i) {
        if (cnt[i]) {
          x |= 1 << i;
          --cnt[i];
        }
      }
      ans = (ans + 1LL * x * x) % mod;
    }
    return ans;
  }
};

```

### Python

```python
class Solution:
    def maxSum(self, nums: List[int], k: int) -> int: mod = 10 ** 9 + 7 cnt = [0] * 31 for x in nums: for i in range(31): if x >> i & 1: cnt[i] += 1 ans = 0 for _ in range(k): x = 0 for i in range(31): if cnt[i]: x |= 1 << i cnt[i] -= 1 ans = (ans + x * x) % mod return ans

```