Apply Operations on Array to Maximize Sum of Squares

Apply Operations on Array to Maximize Sum of Squares - Problem

You are given a 0-indexed integer array nums and a positive integer k.

You can perform the following operation on the array any number of times:

Choose any two distinct indices i and j
Simultaneously update 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.

After performing operations, you must 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 10⁹ + 7.

Input & Output

Example 1 — Basic Case

$ Input: nums = [2,6,4], k = 2

› Output: 72

💡 Note: After operations, we can redistribute bits optimally. With bit counts [0,2,2], we can achieve [0,6,6] by concentrating bits. Taking k=2 largest gives sum = 6² + 6² = 72.

Example 2 — Single Element

$ Input: nums = [4], k = 1

› Output: 16

💡 Note: Only one element, no operations possible. Sum = 4² = 16.

Example 3 — Larger Array

$ Input: nums = [1,2,3,4,5], k = 3

› Output: 109

💡 Note: With nums=[1,2,3,4,5] and k=3, optimal bit distribution concentrates available bits to maximize the sum of squares of the 3 largest values.

Constraints

1 ≤ nums.length ≤ 10⁵
1 ≤ nums[i] ≤ 10⁹
1 ≤ k ≤ nums.length

Visualization

Tap to expand

Asked in

G Google 25 M Microsoft 20 A Amazon 15

The key insight is that AND/OR operations preserve the total number of bits at each position. The optimal strategy is to greedily distribute bits from most significant to least significant positions to maximize the k largest values. Best approach uses bit frequency counting with greedy assignment. Time: O(n log max_val + k log k), Space: O(1).

Common Approaches

✓ Bit Frequency Greedy

⏱️ Time: O(n log max_val + k log k) Space: O(log max_val)

The key insight is that AND/OR operations preserve the total number of 1-bits at each position. We count all bits, then greedily assign them to create the largest possible numbers.

Brute Force with All Operations

⏱️ Time: O(2^(n²) × n log n) Space: O(n)

Generate all possible states by trying every pair of indices for operations, then select k largest elements for each state. This explores the complete solution space but is exponentially slow.

Greedy with Priority Queue

⏱️ Time: O(n log max_val + k log k × log max_val) Space: O(k + log max_val)

Similar to bit frequency approach but uses a priority queue to efficiently manage the k largest values during bit distribution, avoiding repeated sorting.

Bit Frequency Greedy — Algorithm Steps

Count frequency of each bit position across all numbers
Greedily distribute bits to create k largest numbers
Calculate sum of squares modulo 10^9+7

Visualization

Tap to expand

Step-by-Step Walkthrough

Count Bits

Count 1s at each bit position

Distribute

Assign bits greedily from MSB to LSB

Calculate

Sum of squares of k largest values

Code -

solution.c — C

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

#define MOD 1000000007

int compare(const void *a, const void *b) {
    return (*(int*)b - *(int*)a);  // Reverse order for largest first
}

int solution(int *nums, int n, int k) {
    // Count frequency of each bit position
    int bitCount[32] = {0};
    for (int i = 0; i < n; i++) {
        for (int bit = 0; bit < 32; bit++) {
            if (nums[i] & (1 << bit)) {
                bitCount[bit]++;
            }
        }
    }
    
    // Greedily assign bits to create k largest numbers
    int *result = calloc(k, sizeof(int));
    
    // For each bit position from most significant to least
    for (int bit = 31; bit >= 0; bit--) {
        int count = bitCount[bit];
        // Sort to assign bits to largest current values
        qsort(result, k, sizeof(int), compare);
        for (int i = 0; i < count && i < k; i++) {
            result[i] |= (1 << bit);
        }
    }
    
    // Calculate sum of squares
    long long total = 0;
    for (int i = 0; i < k; i++) {
        total = (total + (long long)result[i] * result[i]) % MOD;
    }
    
    free(result);
    return (int)total;
}

int main() {
    char line[1000];
    fgets(line, sizeof(line), stdin);
    
    int nums[100], n = 0;
    char *p = line;
    while (*p && *p != '[') p++;
    if (*p == '[') p++;
    
    while (*p && *p != ']') {
        while (*p && (*p == ' ' || *p == ',')) p++;
        if (*p == ']' || *p == '\0') break;
        nums[n++] = (int)strtol(p, &p, 10);
    }
    
    int k;
    scanf("%d", &k);
    
    printf("%d\n", solution(nums, n, k));
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n log max_val + k log k)

Count bits for each number plus sorting k results

⚠ Quadratic Growth

Space Complexity

O(log max_val)

Array to store bit frequencies (up to 32 bits for integers)

⚡ Linearithmic Space

33.0K Views

Medium Frequency

~35 min Avg. Time

890 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

Constraints

Visualization

Related Problems

Common Approaches

Bit Frequency Greedy — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler