Apply operations on array to maximize sum of squares

Array Hash Table Greedy Bit Manipulation

44.3% acceptance

Feb 25, 2026

196

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 i and j and simultaneously update the values of 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. 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.

Solution

Rust

Time O(n²)

Space O(1)

LeetCode

solution.rs

impl Solution {
  pub fn max_sum(nums: Vec<i32>, k: i32) -> i32 {
    const MOD: i64 = 1_000_000_007;
    // Key insight: AND+OR operations preserve the multiset of bits at each position.
    // We can freely move bits between numbers - so count bit frequencies and greedily
    // construct the k largest numbers by filling highest bits first.
    let mut bit_count = [0i32; 30];
    for &n in &nums {
      for b in 0..30 {
        if n & (1 << b) != 0 {
          bit_count[b] += 1;
        }
      }
    }
    let mut result: i64 = 0;
    let k = k as i32;
    for _ in 0..k {
      let mut val: i64 = 0;
      for b in 0..30 {
        if bit_count[b] > 0 {
          val |= 1 << b;
          bit_count[b] -= 1;
        }
      }
      result = (result + val * val % MOD) % MOD;
    }
    result as i32
  }
}