Apply Operations on Array to Maximize Sum of Squares

Apply Operations on Array to Maximize Sum of Squares - Problem

You're given a 0-indexed integer array nums and a positive integer k. Your goal is to strategically manipulate the array using bitwise operations to maximize the sum of squares of the k largest elements.

Available Operation: Choose any two distinct indices i and j, then simultaneously:

Update nums[i] to nums[i] AND nums[j]
Update nums[j] to nums[i] OR nums[j]

You can perform this operation any number of times. The key insight is that this operation allows you to redistribute bits between array elements - the AND operation concentrates common bits in one element, while the OR operation preserves all bits in the other.

After all operations, select k elements and return the maximum possible sum of their squares modulo 10⁹ + 7.

Example: With nums = [2, 7, 1], k = 2, you can transform the array to maximize two squares by concentrating bits optimally.

Input & Output

example_1.py — Basic Case

$ Input: {"nums": [2, 7, 1], "k": 2}

› Output: 58

💡 Note: We can redistribute bits optimally: count bits at each position (pos 0: 2, pos 1: 2, pos 2: 1), then create two numbers [7, 3] by assigning bits greedily. Sum of squares: 7² + 3² = 49 + 9 = 58.

example_2.py — Single Element

$ Input: {"nums": [1, 2, 3, 4, 5], "k": 3}

› Output: 261

💡 Note: Bit counts: pos 0: 3, pos 1: 2, pos 2: 2. Optimal assignment creates [7, 6, 4]. Sum: 7² + 6² + 4² = 49 + 36 + 16 = 101. Wait, let me recalculate: [1,2,3,4,5] has bits [001,010,011,100,101]. Total bits: pos0: 3, pos1: 2, pos2: 2. Creating [7,6,4]: 49+36+16=101.

example_3.py — Edge Case k=1

$ Input: {"nums": [1, 15, 6], "k": 1}

› Output: 225

💡 Note: With k=1, we can concentrate all bits into one number. Bit counts: pos 0: 2, pos 1: 2, pos 2: 2, pos 3: 1. Creates one optimal number: 15 (binary 1111). Sum: 15² = 225.

Visualization

Tap to expand

Understanding the Visualization

Inventory Count

Count how many coins of each type (bit position) exist across all people

Value Priority

Prioritize higher-value coins (more significant bits) for maximum impact

Greedy Distribution

Give the most valuable coins to your k chosen people first

Optimize Squares

This maximizes the sum of squared wealth across your k millionaires

Key Takeaway

🎯 Key Insight: AND/OR operations preserve total bit counts, so we can optimally redistribute bits by counting each position and greedily assigning them to maximize the sum of k squares.

Time & Space Complexity

Time Complexity

⏱️

O(30n)

30 iterations for each bit position times n elements to count bits

✓ Linear Growth

Space Complexity

O(1)

Only using arrays of size 30 for bit counts and k for result

✓ Linear Space

Constraints

1 ≤ nums.length ≤ 10⁵
1 ≤ nums[i] ≤ 10⁹
1 ≤ k ≤ nums.length
Answer must be returned modulo 10⁹ + 7

Asked in

G Google 42 ⊞ Microsoft 38 a Amazon 31 f Meta 24

The optimal solution uses bit manipulation insights: since AND/OR operations preserve total bit counts at each position, we can count all bits and greedily redistribute them to create k maximum values. Time: O(30n), Space: O(1).

Common Approaches

Approach	Time	Space	Notes
✓ Greedy Bit Redistribution (Optimal)	O(30n)	O(1)	Count bits at each position and optimally redistribute to k largest values

Greedy Bit Redistribution (Optimal) — Algorithm Steps

Count the frequency of each bit position (0-29) across all numbers
For each bit position from most significant to least significant:
Assign as many 1s as possible to the first k positions in our result array
Build k numbers by combining the assigned bits
Return sum of squares of these k numbers modulo 10^9 + 7

Visualization

Tap to expand

Step-by-Step Walkthrough

Count All Bits

Count 1s at each bit position across all numbers

Greedy Assignment

Assign most significant bits first to k positions

Build Optimal Numbers

Construct k numbers with maximum possible values

Calculate Result

Sum squares of the k optimized numbers

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>

#define MOD 1000000007

long long maxSumOfSquares(int *nums, int n, int k) {
    // Count bits at each position
    int bitCount[30] = {0};
    for (int i = 0; i < n; i++) {
        for (int bit = 0; bit < 30; bit++) {
            if (nums[i] & (1 << bit)) {
                bitCount[bit]++;
            }
        }
    }
    
    // Create k optimal numbers
    long long *result = calloc(k, sizeof(long long));
    
    // For each bit position (from most significant)
    for (int bit = 29; bit >= 0; bit--) {
        int count = bitCount[bit];
        // Assign bits greedily
        for (int i = 0; i < count && i < k; i++) {
            result[i] |= (1LL << bit);
        }
    }
    
    // Calculate sum of squares
    long long total = 0;
    for (int i = 0; i < k; i++) {
        total = (total + (result[i] * result[i]) % MOD) % MOD;
    }
    
    free(result);
    return total;
}

int main() {
    int nums[] = {2, 7, 1};
    int k = 2;
    printf("%lld\n", maxSumOfSquares(nums, 3, k));
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(30n)

30 iterations for each bit position times n elements to count bits

✓ Linear Growth

Space Complexity

O(1)

Only using arrays of size 30 for bit counts and k for result

✓ Linear Space

Constraints

1 ≤ nums.length ≤ 10⁵
1 ≤ nums[i] ≤ 10⁹
1 ≤ k ≤ nums.length
Answer must be returned modulo 10⁹ + 7

26.1K Views

Medium-High Frequency

~25 min Avg. Time

892 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Visualization

Time & Space Complexity

Related Problems

Constraints

Common Approaches

Greedy Bit Redistribution (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Constraints

Select Compiler