Apply Operations on Array to Maximize Sum of Squares solution and explanation
6 more videos on this problem
2897. Apply Operations on Array to Maximize Sum of Squares | Weekly Leetcode 366codingMohan · 31:08 · 1K views
2896. Apply Operations on Array to Maximize Sum of Squares (Leetcode Hard)Programming Live with Larry · 13:11 · 453 views
【每日一题】LeetCode 2897. Apply Operations on Array to Maximize Sum of SquaresHuifeng Guan · 14:11 · 366 views
Apply Operations on Array to Maximize Sum of Squares | Weekly-contest-366Rohit Satpute · 12:06 · 254 views
Weekly Contest 366 || 2897. Apply Operations on Array to Maximize Sum of Squares(Hard)Manas Kumar Panda · 28:09 · 174 views
LEETCODE WEEKLY CONTEST 366| ABCD |2897 Apply Operations on Array to Maximize Sum of Squares 8.10.23CodingVerse · 10:54 · 96 views
Apply Operations on Array to Maximize Sum of Squares (LeetCode Hard) (Python)WYCode · 6:02 · 44 views
The 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
iandjand simultaneously update the values ofnums[i]to(nums[i] AND nums[j])andnums[j]to(nums[i] OR nums[j]). Here,ORdenotes the bitwiseORoperation, andANDdenotes the bitwiseANDoperation.
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 <= 1051 <= nums[i] <= 109
Approach overview
Problem Overview: You are given an array of integers and can repeatedly apply an operation that replaces two numbers with their bitwise AND and OR. After performing operations, choose k elements so the sum of their squares is maximized. The key observation: these operations never change the total count of set bits at each bit position across the array.
Approach 1: Greedy Bit Distribution (O((n + k) * B) time, O(B) space)
Track how many times each bit appears across the array. Because the AND/OR operation only redistributes bits between numbers, the total count of every bit position remains constant. To maximize the sum of squares, construct the largest possible numbers greedily. For each of the k selections, iterate through bit positions and assign a bit if any occurrences remain, decrementing the count. This packs high-value bits together, which increases the square of the resulting number. Bit counting takes O(n * B) where B is the number of bits (≈31), and constructing k values takes O(k * B). This approach relies heavily on bit manipulation and a greedy strategy.
Approach 2: Priority Queue (Heap) Simulation (O((n + k) log n) time, O(n) space)
Another way to reason about the redistribution is to dynamically form the largest numbers using available bits and always process the current largest candidates. Maintain numbers in a max heap and repeatedly combine available bit contributions to grow the largest values before squaring them. The heap ensures the largest partially formed values are expanded first, which aligns with maximizing the square contribution. Each heap operation costs O(log n), making the overall complexity O((n + k) log n). This approach is useful if you prefer explicit element tracking using a priority queue instead of direct bit counting.
Recommended for interviews: The greedy bit-counting solution is the expected optimal approach. It demonstrates the critical insight that AND/OR operations conserve bit frequencies, allowing you to ignore the actual operations and work directly with bit counts. Interviewers typically look for this observation first, then the greedy construction that groups bits to maximize squared values.
Approach 1: Greedy Approach to Maximize Final Elements
The goal is to maximize the k largest elements in the array as they substantially increase the sum of squares. By choosing the largest element available and performing bitwise operations wisely, we can pump up the values further.
Initially observe that the AND operation sets insignificant bits to zero, reducing the smallest values to zero over iterations, thus promoting the larger numbers. Use sorting to easily select the largest squares.
def max_sum_of_squares(nums, k):
nums.sort(reverse=True)
mod = 10**9 + 7
return sum((x * x) % mod for x in nums[:k]) % mod
nums = [2, 6, 5, 8]
k = 2
print(max_sum_of_squares(nums, k))import java.util.Arrays;
public class MaxSumOfSquares {
private static final int MOD = 1_000_000_007;
public static long maxSumOfSquares(int[] nums, int k) {
Arrays.sort(nums);
long sum = 0;
for (int i = nums.length - 1; i >= nums.length - k; i--) {
sum = (sum + (long)nums[i] * nums[i]) % MOD;
}
return sum;
}
public static void main(String[] args) {
int[] nums = {2, 6, 5, 8};
int k = 2;
System.out.println(maxSumOfSquares(nums, k));
}
}#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
const int MOD = 1e9 + 7;
long long maxSumOfSquares(vector<int>& nums, int k) {
sort(nums.rbegin(), nums.rend());
long long sum = 0;
for (int i = 0; i < k; ++i) {
sum = (sum + (long long)nums[i] * nums[i]) % MOD;
}
return sum;
}
int main() {
vector<int> nums = {2, 6, 5, 8};
int k = 2;
cout << maxSumOfSquares(nums, k) << endl;
return 0;
}function maxSumOfSquares(nums, k) {
nums.sort((a, b) => b - a);
const MOD = 1e9 + 7;
let sum = 0;
for (let i = 0; i < k; i++) {
sum = (sum + nums[i] * nums[i]) % MOD;
}
return sum;
}
let nums = [2, 6, 5, 8];
let k = 2;
console.log(maxSumOfSquares(nums, k));#include <stdio.h>
#include <stdlib.h>
int cmp(const void *a, const void *b) {
return (*(int*)b - *(int*)a);
}
#define MOD 1000000007
long long maxSumOfSquares(int* nums, int numsSize, int k) {
qsort(nums, numsSize, sizeof(int), cmp);
long long sum = 0;
for (int i = 0; i < k; i++) {
sum = (sum + ((long long)nums[i] * nums[i]) % MOD) % MOD;
}
return sum;
}
int main() {
int nums[] = {2, 6, 5, 8};
int numsSize = sizeof(nums) / sizeof(nums[0]);
int k = 2;
printf("%lld\n", maxSumOfSquares(nums, numsSize, k));
return 0;
}using System;
class MaxSumOfSquares {
const int MOD = 1000000007;
static long MaxSumOfSquares(int[] nums, int k) {
Array.Sort(nums);
long sum = 0;
for (int i = nums.Length - 1; i >= nums.Length - k; i--) {
sum = (sum + (long)nums[i] * nums[i]) % MOD;
}
return sum;
}
static void Main() {
int[] nums = { 2, 6, 5, 8 };
int k = 2;
Console.WriteLine(MaxSumOfSquares(nums, k));
}
}In Python, the list nums is sorted in descending order with the help of the sort() method. The sum of squares of the largest k elements is computed and reduced modulo 10^9 + 7.
Time Complexity: O(n log n), from sorting. Space Complexity: O(1), without additional data structures.
Approach 2: Priority Queue (Heap) Approach to Dynamically Track Largest Elements
This method employs a max heap to dynamically track the k largest elements in the array, which will allow efficient access without sorting the entire array. This is particularly useful for very large arrays where sorting might be inefficient.
import heapq
def max_sum_of_squares(nums, k):
mod = 10**9 + 7
heap = [-x for x in nums] # create a max heap
heapq.heapify(heap)
res = 0
for _ in range(k):
largest = -heapq.heappop(heap)
res = (res + largest * largest) % mod
return res
nums = [2, 6, 5, 8]
k = 2
print(max_sum_of_squares(nums, k))import java.util.PriorityQueue;
public class MaxSumOfSquares {
private static final int MOD = 1_000_000_007;
public static long maxSumOfSquares(int[] nums, int k) {
PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);
for (int num : nums) {
maxHeap.add(num);
}
long sum = 0;
for (int i = 0; i < k; i++) {
int largest = maxHeap.poll();
sum = (sum + (long)largest * largest) % MOD;
}
return sum;
}
public static void main(String[] args) {
int[] nums = {2, 6, 5, 8};
int k = 2;
System.out.println(maxSumOfSquares(nums, k));
}
}#include <iostream>
#include <vector>
#include <queue>
using namespace std;
const int MOD = 1e9 + 7;
long long maxSumOfSquares(vector<int>& nums, int k) {
priority_queue<int> maxHeap(nums.begin(), nums.end());
long long sum = 0;
while (k-- > 0) {
int top = maxHeap.top();
maxHeap.pop();
sum = (sum + (long long)top * top) % MOD;
}
return sum;
}
int main() {
vector<int> nums = {2, 6, 5, 8};
int k = 2;
cout << maxSumOfSquares(nums, k) << endl;
return 0;
}class MaxHeap {
constructor() {
this.heap = [];
}
insert(val) {
this.heap.push(val);
this.heapifyUp();
}
extractMax() {
const max = this.heap[0];
const end = this.heap.pop();
if (this.heap.length > 0) {
this.heap[0] = end;
this.heapifyDown();
}
return max;
}
heapifyUp() {
let index = this.heap.length - 1;
const element = this.heap[index];
while (index > 0) {
let parentIndex = Math.floor((index - 1) / 2);
let parent = this.heap[parentIndex];
if (element <= parent) break;
this.heap[index] = parent;
index = parentIndex;
}
this.heap[index] = element;
}
heapifyDown() {
let index = 0;
const length = this.heap.length;
const element = this.heap[0];
while (true) {
let leftChildIndex = 2 * index + 1;
let rightChildIndex = 2 * index + 2;
let leftChild, rightChild;
let swap = null;
if (leftChildIndex < length) {
leftChild = this.heap[leftChildIndex];
if (leftChild > element) swap = leftChildIndex;
}
if (rightChildIndex < length) {
rightChild = this.heap[rightChildIndex];
if (
(swap === null && rightChild > element) ||
(swap !== null && rightChild > leftChild)
) swap = rightChildIndex;
}
if (swap === null) break;
this.heap[index] = this.heap[swap];
index = swap;
}
this.heap[index] = element;
}
}
function maxSumOfSquares(nums, k) {
const maxHeap = new MaxHeap();
for (const num of nums) {
maxHeap.insert(num);
}
const MOD = 1e9 + 7;
let sum = 0;
for (let i = 0; i < k; i++) {
const largest = maxHeap.extractMax();
sum = (sum + largest * largest) % MOD;
}
return sum;
}
let nums = [2, 6, 5, 8];
let k = 2;
console.log(maxSumOfSquares(nums, k));#include <stdio.h>
#include <stdlib.h>
#include <limits.h>
#define MOD 1000000007
void swap(int* a, int* b) {
int temp = *a;
*a = *b;
*b = temp;
}
void heapify(int arr[], int n, int i) {
int largest = i;
int left = 2 * i + 1;
int right = 2 * i + 2;
if (left < n && arr[left] > arr[largest])
largest = left;
if (right < n && arr[right] > arr[largest])
largest = right;
if (largest != i) {
swap(&arr[i], &arr[largest]);
heapify(arr, n, largest);
}
}
long long maxSumOfSquares(int* nums, int numsSize, int k) {
for (int i = numsSize / 2 - 1; i >= 0; i--)
heapify(nums, numsSize, i);
long long sum = 0;
for (int i = 0; i < k; i++) {
sum = (sum + ((long long)nums[0] * nums[0]) % MOD) % MOD;
nums[0] = nums[numsSize - 1];
numsSize--;
heapify(nums, numsSize, 0);
}
return sum;
}
int main() {
int nums[] = {2, 6, 5, 8};
int numsSize = sizeof(nums) / sizeof(nums[0]);
int k = 2;
printf("%lld\n", maxSumOfSquares(nums, numsSize, k));
return 0;
}using System;
using System.Collections.Generic;
class MaxSumOfSquares {
const int MOD = 1000000007;
static long MaxSumOfSquares(int[] nums, int k) {
PriorityQueue<int, int> maxHeap = new(nums.Select(x => (x, -x)));
long sum = 0;
for (int i = 0; i < k; i++) {
int largest = maxHeap.Dequeue();
sum = (sum + (long)largest * largest) % MOD;
}
return sum;
}
static void Main() {
int[] nums = { 2, 6, 5, 8 };
int k = 2;
Console.WriteLine(MaxSumOfSquares(nums, k));
}
}This Python code uses a max heap (simulated with a min heap using negative values) to extract the largest k elements. This ensures efficient computation of the sum of squares of the largest elements.
Time Complexity: O(n + k log n) with heap operations. Space Complexity: O(n) due to heap size.
Approach 3: Bitwise Operation + Greedy
According to the problem description, for an operation, we can change nums[i] to nums[i] AND nums[j], and change nums[j] to nums[i] OR nums[j]. Let's consider the bits of the numbers. If two bits are both 1 or both 0, the result of the operation will not change the bits. If two bits are different, the result of the operation will change the bits to 0 and 1, respectively. Therefore, we can move 1 bits to 0 bits, but not vice versa.
We can use an array cnt to count the number of 1 bits in each position, and then select k numbers from them. To maximize the sum of squares, we should choose the largest numbers as much as possible. This is because, assuming the sum of squares of two numbers is a^2 + b^2 (where a \gt b), changing them to (a + c)^2 + (b - c)^2 = a^2 + b^2 + 2c(a - b) + 2c^2 \gt a^2 + b^2 will increase the sum of squares. Therefore, to maximize the sum of squares, we should choose the largest number.
The time complexity is O(n times log M), and the space complexity is O(log M). Here, M is the maximum value in the array.
function maxSum(nums: number[], k: number): number {
const cnt: number[] = Array(31).fill(0);
for (const x of nums) {
for (let i = 0; i < 31; ++i) {
if ((x >> i) & 1) {
++cnt[i];
}
}
}
let ans = 0n;
const mod = 1e9 + 7;
while (k-- > 0) {
let x = 0;
for (let i = 0; i < 31; ++i) {
if (cnt[i] > 0) {
x |= 1 << i;
--cnt[i];
}
}
ans = (ans + BigInt(x) * BigInt(x)) % BigInt(mod);
}
return Number(ans);
}func maxSum(nums []int, k int) (ans int) {
cnt := [31]int{}
for _, x := range nums {
for i := 0; i < 31; i++ {
if x>>i&1 == 1 {
cnt[i]++
}
}
}
const mod int = 1e9 + 7
for ; k > 0; k-- {
x := 0
for i := 0; i < 31; i++ {
if cnt[i] > 0 {
x |= 1 << i
cnt[i]--
}
}
ans = (ans + x*x) % mod
}
return
}Complexity comparison
| Approach | Complexity |
|---|---|
| Greedy Approach to Maximize Final Elements | Time Complexity: O(n log n), from sorting. Space Complexity: O(1), without additional data structures. |
| Priority Queue (Heap) Approach to Dynamically Track Largest Elements | Time Complexity: O(n + k log n) with heap operations. Space Complexity: O(n) due to heap size. |
| Bitwise Operation + Greedy | – |
Complexity analysis
| Approach | Time | Space | When to Use |
|---|---|---|---|
| Greedy Bit Distribution | O((n + k) * B) | O(B) | Best overall approach when bit limits are small (≈31). Most efficient and common interview solution. |
| Priority Queue (Heap) Simulation | O((n + k) log n) | O(n) | Useful when explicitly tracking element growth or when modeling redistribution dynamically. |
Frequently asked questions
Is Apply Operations on Array to Maximize Sum of Squares easy or hard?
The problem is rated Hard because the optimal solution requires recognizing that AND and OR operations preserve bit counts. Once this invariant is identified, the greedy bit construction becomes straightforward, but discovering the insight is the challenging part.
Apply Operations on Array to Maximize Sum of Squares Python/Java solution
Most implementations follow the same structure across languages: count set bits for each position, build k numbers greedily by consuming bits, and accumulate the squared values with modulo arithmetic. This approach works efficiently in Python, Java, C++, C#, and JavaScript.
How to solve Apply Operations on Array to Maximize Sum of Squares in O(n)?
Strict O(n) is not typical because the algorithm iterates over bit positions. The closest practical complexity is O((n + k) * B). First count how many times each bit appears, then greedily construct k numbers by assigning available bits from highest contribution. Each constructed number is squared and added to the result.
What is the best approach for Apply Operations on Array to Maximize Sum of Squares?
The optimal approach is a greedy bit-counting strategy. Count how many times each bit appears in the array, then greedily construct the largest possible numbers for k iterations by consuming available bits. This maximizes the squared contribution of each number. The complexity is O((n + k) * B) where B is the number of bits (around 31).
Is Apply Operations on Array to Maximize Sum of Squares asked at Google/Amazon/Meta?
Problems involving greedy bit redistribution and maximizing squared sums appear frequently in interviews at companies like Google, Amazon, and Meta. Variants often test understanding of bit manipulation, greedy construction, and mathematical reasoning about bit contributions.
What data structure is used in Apply Operations on Array to Maximize Sum of Squares?
The main solution uses an array to store bit frequencies and applies greedy construction using bit manipulation. An alternative approach can use a priority queue (max heap) to dynamically track and expand the largest numbers during construction.
What is the time complexity of Apply Operations on Array to Maximize Sum of Squares?
The optimal greedy solution runs in O((n + k) * B) time, where n is the array size and B is the number of bit positions (typically 31 for integers). Counting bits takes O(n * B), and building k values takes O(k * B). Space complexity is O(B).