Skip to content

2897. Apply Operations on Array to Maximize Sum of Squares 👍

  • Time: $O(30n) = O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
 public:
  int maxSum(vector<int>& nums, int k) {
    constexpr int kMod = 1'000'000'007;
    constexpr int kMaxBit = 30;
    int ans = 0;
    // minIndices[i] := the minimum index in `optimalNums` that the i-th bit
    // should be moved to
    vector<int> minIndices(kMaxBit);
    vector<int> optimalNums(nums.size());

    for (const int num : nums)
      for (int i = 0; i < kMaxBit; ++i)
        if (num >> i & 1)
          optimalNums[minIndices[i]++] |= 1 << i;

    for (int i = 0; i < k; ++i)
      ans = (ans + static_cast<long>(optimalNums[i]) * optimalNums[i]) % kMod;

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
  public int maxSum(List<Integer> nums, int k) {
    final int kMod = 1_000_000_007;
    final int kMaxBit = 30;
    int ans = 0;
    // minIndices[i] := the minimum index in `optimalNums` that the i-th bit
    // should be moved to
    int[] minIndices = new int[kMaxBit];
    int[] optimalNums = new int[nums.size()];

    for (final int num : nums)
      for (int i = 0; i < kMaxBit; i++)
        if ((num >> i & 1) == 1)
          optimalNums[minIndices[i]++] |= 1 << i;

    for (int i = 0; i < k; i++)
      ans = (int) (((long) ans + (long) optimalNums[i] * optimalNums[i]) % kMod);

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
  def maxSum(self, nums: list[int], k: int) -> int:
    kMod = 1_000_000_007
    kMaxBit = 30
    ans = 0
    # minIndices[i] := the minimum index in `optimalNums` that the i-th bit
    # should be moved to
    minIndices = [0] * kMaxBit
    optimalNums = [0] * len(nums)

    for num in nums:
      for i in range(kMaxBit):
        if num >> i & 1:
          optimalNums[minIndices[i]] |= 1 << i
          minIndices[i] += 1

    for i in range(k):
      ans += optimalNums[i]**2
      ans %= kMod

    return ans