LeetCode 2897 - Apply Operations on Array to Maximize Sum of Squares

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Greedy, Bit Manipulation

Solution

Problem Understanding

The problem presents an integer array nums and a positive integer k. You can perform a bitwise operation on any two distinct elements of the array any number of times. Specifically, for indices i and j, you update nums[i] to (nums[i] AND nums[j]) and nums[j] to (nums[i] OR nums[j]). After applying such operations, the goal is to select exactly k elements from the array and maximize the sum of their squares, returning the result modulo 10^9 + 7.

Key points to note are:

The array can be large, up to 10^5 elements.
Each element can be up to 10^9, meaning up to 30 bits.
The AND and OR operations follow standard bitwise rules.
The objective is not the final array itself but maximizing the sum of squares of the largest k numbers achievable.

A naive approach might attempt all possible pairwise operations, but this is infeasible because the number of possible operation sequences grows extremely quickly. Therefore, the challenge lies in understanding the bitwise operation behavior and leveraging it to compute the maximal values without brute-forcing every combination.

Important edge cases include:

k equals the length of the array, meaning we must sum squares of all elements.
nums already contains very large values; applying operations might reduce smaller numbers unnecessarily.
All numbers are equal; operations may have no effect.

Approaches

Brute Force

A brute-force approach would involve simulating all possible pairwise operations. For each pair (i, j), apply the operation, recurse on the resulting array, and track the maximum sum of squares for any selection of k elements. This guarantees correctness because it explores all possibilities, but its complexity is exponential:

Time complexity: O(2^(n^2)) in the worst case, since each operation changes the array and leads to many branches.
Space complexity: O(n) for recursion depth plus the array copy storage.

Clearly, this approach is infeasible for n up to 10^5.

Optimal Approach

The key observation is that for any two numbers a and b:

a AND b <= a and a AND b <= b
a OR b >= a and a OR b >= b

This implies the OR operation increases a number or keeps it the same, while AND decreases a number or keeps it the same. Since we want the largest k numbers, the optimal strategy is to apply operations that move bits from smaller numbers to larger numbers.

Effectively, the operation can be used to accumulate all set bits into a single number, which will become the largest possible element. After that, the remaining elements are smaller or zero. The algorithm becomes:

Compute the bitwise OR of all numbers to get the largest achievable number.
Count this OR number once and combine it with the largest k-1 other original numbers (or zeros if fewer exist) to maximize the sum of squares.

This is efficient because it avoids simulating every operation and directly calculates the maximal values.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(2^(n^2))	O(n)	Simulates all operation sequences; infeasible
Optimal	O(n log n)	O(n)	Use bitwise OR to find largest possible number, then select largest k elements

Algorithm Walkthrough

Initialize a variable max_val to 0. This will store the bitwise OR of all elements.
Iterate through each number in nums, updating max_val with max_val |= num. After this loop, max_val represents the largest possible number achievable through operations.
Sort the array nums in descending order to easily select the largest k numbers.
Replace the first element with max_val because we can accumulate all bits into a single element.
Select the first k elements of the sorted array (after updating the first element) as these are the numbers that maximize the sum of squares.
Compute the sum of squares of these k numbers, taking modulo 10^9 + 7.

Why it works: The bitwise OR property ensures that all possible bits can be accumulated in a single number. No operation can increase any number beyond this OR. Sorting guarantees that the largest numbers are chosen for squaring. Therefore, this procedure guarantees the maximum sum of squares.

Python Solution

class Solution:
    def maxSum(self, nums: list[int], k: int) -> int:
        MOD = 10**9 + 7
        max_val = 0
        for num in nums:
            max_val |= num  # accumulate all bits into one number
        
        nums.sort(reverse=True)
        nums[0] = max_val  # ensure the largest number is maximized
        
        result = sum(x * x for x in nums[:k]) % MOD
        return result

Implementation Explanation:

The for loop computes the overall OR of all numbers.
Sorting in descending order makes it trivial to pick the largest k numbers.
Replacing nums[0] with max_val ensures the largest number is as large as possible.
The sum of squares modulo 10^9 + 7 is computed in a concise generator expression.

Go Solution

import "sort"

func maxSum(nums []int, k int) int {
    const MOD = 1_000_000_007
    maxVal := 0
    for _, num := range nums {
        maxVal |= num
    }
    
    sort.Slice(nums, func(i, j int) bool {
        return nums[i] > nums[j]
    })
    
    nums[0] = maxVal
    result := 0
    for i := 0; i < k; i++ {
        result = (result + nums[i]*nums[i]) % MOD
    }
    return result
}

Go Implementation Notes:

Uses sort.Slice with a comparator to sort in descending order.
Integer overflow is handled by taking modulo at each addition.
Otherwise, the algorithm mirrors the Python version.

Worked Examples

Example 1

nums = [2,6,5,8], k = 2

Compute max_val = 2 | 6 | 5 | 8 = 15.
Sort nums descending: [8,6,5,2].
Replace nums[0] with 15: [15,6,5,2].
Take first k=2 elements: [15,6].
Sum of squares: 15^2 + 6^2 = 225 + 36 = 261.

Example 2

nums = [4,5,4,7], k = 3

Compute max_val = 4 | 5 | 4 | 7 = 7.
Sort nums descending: [7,5,4,4].
Replace nums[0] with 7: [7,5,4,4].
Take first k=3 elements: [7,5,4].
Sum of squares: 49 + 25 + 16 = 90.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log n)	Sorting dominates the O(n) OR computation
Space	O(1)	In-place sorting and accumulation; no extra structures

Sorting is the main bottleneck, while computing the bitwise OR is linear. Memory is minimal, making this approach suitable for large arrays.

Test Cases

# Provided examples
assert Solution().maxSum([2,6,5,8], 2) == 261  # Max sum with operations
assert Solution().maxSum([4,5,4,7], 3) == 90   # No operations needed

# Edge cases
assert Solution().maxSum([1], 1) == 1  # Single element
assert Solution().maxSum([1,2,3], 3) == 14  # Sum of squares of all elements after OR
assert Solution().maxSum([1,1,1,1], 2) == 2  # All identical elements
assert Solution().maxSum([10**9, 1, 2], 1) == 10**18 % (10**9 + 7)  # Large numbers

Test	Why
`[2,6,5,8], 2`	Validates accumulation of bits and selecting top k
`[4,5,4,7], 3`	Validates scenario where no operations are needed
`[1],1`	Minimal input case
`[1,2,3],3`	Tests selecting all elements after bit accumulation
`[1,1,1,1],2`	All elements identical, ensures correct sum of squares
`[10^9,1,2],1`	Tests large numbers and modulo handling

Edge Cases

Single-element array: If nums has only one element, the OR operation is trivial. The implementation correctly