2897. Apply Operations on Array to Maximize Sum of Squares
Problem Description
In the given problem, we have an array of integers referred to as nums
, and a positive integer k
. The task is to perform operations on the elements of this array to ultimately select k
elements from it and maximize the sum of their squares.
The operation that can be performed involves choosing any two distinct indices i
and j
in the array and replacing nums[i]
with the result of the bitwise AND (&
) of nums[i]
and nums[j]
, while nums[j]
is replaced with the result of the bitwise OR (|
) of the same elements. The important thing to notice here is that the AND operation tends to reduce the numbers by turning 1s to 0s when the corresponding bits are different, while the OR operation does the opposite. However, once a bit is turned to 0 by the AND operation, it cannot be reverted back to 1.
After performing these operations any number of times, we are supposed to take k
numbers from the resulting array and calculate the sum of their squares. The goal is to find out how to perform these operations such that this sum is maximized. Finally, the sum of squares should be returned modulo 10^9 + 7
, due to the possibility of the sum becoming very large.
Intuition
Understanding how bitwise operations work is crucial to solving this problem. When performing the AND operation between two numbers, the bits that are different (where one is 1 and the other is 0) will result in 0. Conversely, the OR operation will produce a 1 bit in places where at least one of the two bits is 1.
The strategy to maximize the sum of squares is to have the numbers with the highest set bits included in the final selection of k
elements, because these will give larger squares due to the exponential growth of value with each additional higher-order bit. This means that if a bit is set at a higher position, it contributes more significantly to the value of a number than a bit set at a lower position.
To systematically achieve this, we can track the frequency of set bits at each bit position in all numbers of the array. We then construct the k
largest numbers by taking the most significant set bit (if available) from the bit frequency array, and then proceed to the next most significant bit, fulfilling the count for k
numbers. With each selection, we decrement the count of set bits because we have utilized that set bit. Finally, we square these numbers and keep a cumulative sum, which is returned modulo 10^9 + 7
.
By doing this, we ensure we have the largest possible numbers that can be formed after the allowed operations, hence maximizing the sum of their squares.
Learn more about Greedy patterns.
Solution Approach
The solution implements a greedy algorithm, combined with bit manipulations, to construct the largest possible numbers from the available set bits.
First, we initialize a counter array cnt
of size 31, which corresponds to bits from 0
to 30
, taking into account that the maximum number within 32-bit signed integer range would not exceed 2^31. This array is used to count how many times each bit is set among all the numbers in the input nums
.
We iterate over every number x
in the input array and for each bit i
from 0
to 30
(since we are using 0-indexed bit positions), we check if the i
-th bit of x
is set by shifting x
to the right by i
positions and using bitwise AND with 1
. If the result is 1
, it means the bit is set, and we increment the count of i
-th bit in cnt
.
After preparing the cnt
array, we are ready to select k
largest numbers. We do this k
times:
- Initialize
x
to0
before each of thek
selections. Thisx
will be the construction of one of thek
largest possible numbers. - Iterate through the bit positions from the most significant to the least significant (
30
to0
), and check if we have any remaining set bits stored incnt
at the current positioni
. If we do, we set this bit inx
(x |= 1 << i
) meaning we contribute this bit to make our number larger. We then decrement the count of set bits at this position since we've used one. - Once we have constructed the number
x
, we add its square toans
keeping it modulated withmod = 10**9 + 7
to handle large numbers and prevent integer overflow issues.
By the end of the k
iterations, we have accumulated the maximum sum of squares possible under the given constraints and operations.
In summary, the data structures used in this algorithm are:
- An array
cnt
to keep track of the count of set bit frequencies. - Variables
x
andans
to store the current number being constructed and the final answer, respectively.
This algorithm efficiently utilizes bit manipulations and a greedy approach to solve the problem in a way that scales well with larger inputs.
Ready to land your dream job?
Unlock your dream job with a 2-minute evaluator for a personalized learning plan!
Start EvaluatorExample Walkthrough
Let's use a small example to illustrate the solution approach.
Suppose we have an array nums = [3, 8, 2]
and k = 2
. The binary representations of nums
are [11, 1000, 10]
, where we can see that 8
has the highest bit set. Our goal is to maximize the sum of squares of any two numbers after performing the described operations.
Following the solution approach:
- We initialize an array
cnt
of size31
to track the frequency of set bits at each position. After iterating overnums
,cnt
would look like this (showing only bit positions 0 to 3 for brevity, with position 3 being the most significant bit):
Position: 0 1 2 3 Count: 1 1 0 1
This indicates that we have one 1
at positions 0
, 1
, and 3
across all numbers.
-
Now we construct the largest number possible. With
k = 2
, we do this twice:- For the first number, we start with our
x = 0
and scan from bit 30 to 0. Reaching position 3, we seecnt[3] = 1
, so we can set this bit on ourx
yieldingx = 1000
. We decrementcnt[3]
by 1 since we've used this set bit. - For the second number, we do the same but, since we've used the set bit at position 3, it's no longer available. Thus, we start with
x = 0
, reach position 1, set it sincecnt[1] = 1
, and nowx = 10
.
- For the first number, we start with our
-
We have two numbers
8
(binary1000
) and2
(binary10
), and we square them and add them together keeping modulo10^9 + 7
.
8^2 + 2^2 = 64 + 4 = 68
The maximum sum of squares we can get is 68
. However, considering the operation described, we would realize that no operation can increase the value of 8
since it's the maximum, and performing an AND operation with 2
or 3
would result in a number smaller than or equal to 2
. Hence, our starting array nums = [3, 8, 2]
is already optimized for selecting k = 2
elements to maximize the sum of their squares.
By applying this approach to the actual problem, we'd be able to find the maximum sum of squares of k
numbers from a larger array, with more complex combinations of set bits, while ensuring we're using a greedy method to obtain the largest numbers possible after performing the allowed operations.
Solution Implementation
1class Solution:
2 def maxSum(self, nums: List[int], k: int) -> int:
3 # Define the modulus variable for large number handling as per the problem statement
4 mod = 10**9 + 7
5
6 # Initialize a list to count the number of set bits at each bit position
7 bit_count = [0] * 31 # Assuming 32-bit integers, ignoring the sign bit
8
9 # Count the number of set bits for each bit position in all numbers
10 for num in nums:
11 for i in range(31): # Traverse through all the bits
12 if num >> i & 1: # Check if the ith bit is set
13 bit_count[i] += 1 # Increment count for this bit position
14
15 # Initialize variable to store the sum of squares of chosen numbers
16 max_sum = 0
17
18 # Loop to choose 'k' numbers
19 for _ in range(k):
20 chosen_num = 0 # This will store the current chosen number with maximum set bits
21
22 # Create the maximum number by setting the highest bits where possible
23 for i in range(31):
24 if bit_count[i]: # If there are bits to set at position i
25 chosen_num |= 1 << i # Set the bit at position i
26 bit_count[i] -= 1 # Decrement count as this bit is now used up
27
28 # Add the square of the chosen number to max_sum and take modulo
29 max_sum = (max_sum + chosen_num * chosen_num) % mod
30
31 # Return the computed max_sum
32 return max_sum
33
1class Solution {
2 public int maxSum(List<Integer> nums, int numIterations) {
3 final int MODULO = (int) 1e9 + 7; // Define the modulo to prevent integer overflow.
4 int[] bitCounts = new int[31]; // Array to keep track of the count of 1 bits at each position.
5
6 // Count the number of 1 bits at each position for all numbers in nums.
7 for (int num : nums) {
8 for (int i = 0; i < 31; ++i) {
9 if ((num >> i & 1) == 1) {
10 ++bitCounts[i];
11 }
12 }
13 }
14
15 long answer = 0; // Initialize the answer which will accumulate the value.
16
17 // Perform calculations for 'numIterations' iterations.
18 while (numIterations-- > 0) {
19 int x = 0; // Initialize x which will hold the current max value.
20
21 // Create x by setting bit i if there is any number with a 1 bit at position i.
22 for (int i = 0; i < 31; ++i) {
23 if (bitCounts[i] > 0) {
24 x |= 1 << i; // Set bit i.
25 --bitCounts[i]; // Decrement the count for bit i.
26 }
27 }
28
29 // Add square of the created number x to the answer.
30 answer = (answer + (long) x * x) % MODULO;
31 }
32
33 // Cast the final answer to int and return it.
34 return (int) answer;
35 }
36}
37
1class Solution {
2public:
3 // Function to calculate the maximum sum of squares of k integers
4 // constructed by bitwise OR of different subsets.
5 int maxSum(vector<int>& nums, int k) {
6 // Initialize a counter array to store the number of 1's at each bit position.
7 int bitCount[31] = {};
8
9 // Count the number of 1's at each bit position for all numbers.
10 for (int num : nums) {
11 for (int bit = 0; bit < 31; ++bit) {
12 if (num & (1 << bit)) {
13 ++bitCount[bit];
14 }
15 }
16 }
17
18 // Initialize the answer as a long long to avoid potential overflow.
19 long long answer = 0;
20 // Constant mod to take answer modulo 1e9+7 for large sums.
21 const int mod = 1e9 + 7;
22
23 // Loop k times to construct k integers.
24 while (k--) {
25 // Variable to store the currently constructed integer using bitwise OR.
26 int currentInt = 0;
27 // For each bit position from 0 to 30, we check if we still have numbers left
28 // with this bit set to 1. If so, we include this bit in our current integer.
29 for (int bit = 0; bit < 31; ++bit) {
30 if (bitCount[bit]) {
31 // Set the bit position in the current integer.
32 currentInt |= (1 << bit);
33 // Decrement the count of available 1's at this bit position.
34 --bitCount[bit];
35 }
36 }
37 // Add the square of the current integer to 'answer' and take modulo.
38 answer = (answer + static_cast<long long>(currentInt) * currentInt) % mod;
39 }
40
41 // Cast the result back to int as per function return type.
42 return static_cast<int>(answer);
43 }
44};
45
1function maxSum(nums: number[], k: number): number {
2 // Initialize a count array with 31 zeros, one for each bit position (assuming 32-bit integers)
3 const bitCounts: number[] = Array(31).fill(0);
4
5 // Count the number of 1s at each bit position for all numbers in the array
6 for (const num of nums) {
7 for (let i = 0; i < 31; ++i) {
8 if ((num >> i) & 1) {
9 bitCounts[i]++;
10 }
11 }
12 }
13
14 // Initialize the answer as a BigInt to support potential large numbers resulting from squaring
15 let answer = 0n;
16
17 // Define modulo constant to prevent overflow
18 const mod = BigInt(1e9 + 7);
19
20 // Loop to perform calculation k times
21 while (k-- > 0) {
22 let constructedNumber = 0;
23
24 // Construct the largest number possible using the available bit counts
25 for (let i = 0; i < 31; ++i) {
26 if (bitCounts[i] > 0) {
27 constructedNumber |= 1 << i; // Set the ith bit
28 bitCounts[i]--; // Decrement the bit count at ith position
29 }
30 }
31
32 // Add the square of the constructed number to the answer, and apply modulo operation
33 answer = (answer + BigInt(constructedNumber) * BigInt(constructedNumber)) % mod;
34 }
35
36 // Return the answer as a number after casting from BigInt
37 return Number(answer);
38}
39
Time and Space Complexity
Time Complexity
The given code consists of two main parts which contribute to the time complexity:
-
The first
for
loop iterates through all the elements innums
and for each element, it loops a constant 31 times to count the number of 1s at each bit position. Suppose there aren
elements in the inputnums
. The time complexity for this part isO(n * 31)
. Since 31 is a constant, we simplify this toO(n)
. -
The second
for
loop runsk
times, wherek
is an input parameter that does not depend on the size ofnums
. Inside thisk
-loop, we have another loop that runs 31 times for the bit counting similar to the first part. The time complexity for this part isO(k * 31)
. Again, since 31 is a constant, this is reduced toO(k)
.
Since k
is independent of n
, and both loops do not nest but run in sequence, we sum their complexities getting O(n) + O(k)
. This indicates that as n
or k
grows, the runtime will increase linearly with whichever is larger. However, if we consider the value of the elements in nums
, specifically M
, the maximum value, we recognize that the bit-operations like x >> i & 1
depend on the maximum number of bits used to represent the numbers in the array. This introduces the term log M
, and therefore the overall time complexity considering the bit-manipulations would be O(n * log M)
.
Space Complexity
For space complexity:
-
We only need extra space for the
cnt
array which stores the bit counts for each of the 31 possible bit positions (assuming 32-bit integers). The size of this array is constant hence contributingO(31)
, which simplifies toO(1)
space complexity. -
Besides the
cnt
array, the code uses a fixed number of integer variables. These are also constant space and do not scale with input sizen
or the valuek
.
Given that these values do not depend on the size of the input array or the elements within, the overall space complexity of the code is O(1)
.
In conjunction with the reference answer, the time complexity is represented as O(n * log M)
and the space complexity as O(log M)
, but as analyzed, the correct space complexity should be O(1)
since the only non-constant space used is the cnt
array which has a fixed size regardless of the input.
Learn more about how to find time and space complexity quickly using problem constraints.
In a binary min heap, the maximum element can be found in:
Recommended Readings
Greedy Introduction div class responsive iframe iframe src https www youtube com embed WTslqPbj7I title YouTube video player frameborder 0 allow accelerometer autoplay clipboard write encrypted media gyroscope picture in picture web share allowfullscreen iframe div When do we use greedy Greedy algorithms tend to solve optimization problems Typically they will ask you to calculate the max min of some value Commonly you may see this phrased in the problem as max min longest shortest largest smallest etc These keywords can be identified by just scanning
LeetCode Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way we
Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https algomonster s3 us east 2 amazonaws com recursion jpg You first
Want a Structured Path to Master System Design Too? Donāt Miss This!