2897. Apply Operations on Array to Maximize Sum of Squares

2026, Oct 11 2 mins read

1. Description

Given an integer array nums and a positive integer k, you may repeatedly select two different positions and replace their values a and b with a AND b and a OR b, respectively. Both bitwise results use the values before the operation.
After any number of operations, including zero, select k elements. Maximize the sum of their squares and return that maximum modulo 10$^9$ + 7.

2. Example

Example 1

Input: nums = [2,6,5,8], k = 2
Output: 261
Explanation: Combining 2 with 8 gives [0,6,5,10]. Combining 5 with 10 then gives [0,6,0,15]. Select 15 and 6 to obtain 15$^2$ + 6$^2$ = 261.

Example 2

Input: nums = [4,5,4,7], k = 3
Output: 90
Explanation: Select 7, 5, and 4 without changing the array, giving 7$^2$ + 5$^2$ + 4$^2$ = 90.

3. Constraints

1 <= k <= nums.length <= 10$^5$
1 <= nums[i] <= 10$^9$

4. Solutions

Greedy with Bit Counting

n = nums.size()
Time complexity: O(n)
Space complexity: O(1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
// AND/OR preserves the count of set bits at each position.
// Concentrate available bits into each of the k numbers to maximize their squared sum.
class Solution {
public:
    int maxSum(vector<int> &nums, int k) {
        const long long MOD = 1e9 + 7;

        array<int, 31> bit_count{};
        for (int num : nums) {
            for (int bit = 0; bit < 31; ++bit) {
                if (num & (1 << bit)) {
                    ++bit_count[bit];
                }
            }
        }

        long long answer = 0;
        for (int i = 0; i < k; ++i) {
            long long num = 0;
            for (int bit = 0; bit < 31; ++bit) {
                if (bit_count[bit] > 0) {
                    num |= 1LL << bit;
                    --bit_count[bit];
                }
            }

            answer = (answer + num * num) % MOD;
        }

        return answer;
    }
};

Search

Tags