DDSA Solutions

2897. Apply Operations on Array to Maximize Sum of Squares

HardView on LeetCode
ArrayHash TableGreedyBit Manipulation
Time: O(n log M + k log M)
Space: O(log M)

Problem Overview

The operation replaces two numbers with their AND and their OR.

Intuition

The operation replaces two numbers with their AND and their OR. Looking at one bit at a time, it can move a set bit from one number to the other, but it never creates or destroys one, so the count of each bit across the array stays fixed. Squaring rewards big numbers, so the best plan piles as many bits as possible into the same few numbers. If you build the largest number first, then the next largest, and so on, number j has bit b exactly when more than j inputs had bit b. That means 30 counters describe every number you could build.

Algorithm

  1. 1Count how many inputs have each of the 30 bits set, visiting only the set bits of each number.
  2. 2Track the largest count. Built numbers past that index are 0 and add nothing.
  3. 3For each j from 0 up to the smaller of k and the largest count, set bit b in the value when its count is more than j.
  4. 4Add the square of that value to the answer modulo 1e9+7, using a 64-bit integer for the product.
  5. 5Return the answer.

Example Walkthrough

Input: nums = [2,6,5,8], k = 2

  1. 1.Bit 0 appears once, bit 1 twice, bit 2 twice, and bit 3 once.
  2. 2.The first built number takes all four bits and is 15.
  3. 3.The second takes bits 1 and 2, which still have a count left, and is 6.
  4. 4.15 squared plus 6 squared is 225 + 36.

Output: 261

Common Pitfalls

  • •The values can reach 1e9, so square them in a 64-bit integer before taking the modulus.
  • •Do not try to simulate the operation on pairs. Counting bits gives the final numbers directly.
  • •Choosing the k largest original values is wrong, because the operation can make numbers larger than any input.
  • •Storing all n built numbers works, but the 30 bit counts already describe each one.
2897.cs
C#
// Approach: The operation moves set bits between numbers without changing how
// many times each bit appears. Squares reward concentration, so greedily pack
// bits into as few numbers as possible. Count each bit; the j-th built number
// (0-indexed) has bit b exactly when more than j inputs had it. Sum the
// squares of the first k built numbers, stopping early once they become 0.
// Time: O(n log M + k log M) Space: O(log M)

public class Solution
{
    public int MaxSum(IList<int> nums, int k)
    {
        const int MOD = 1_000_000_007;
        const int BITS = 30;
        int[] count = new int[BITS];
        int maxCount = 0;

        foreach (int num in nums)
        {
            uint x = (uint)num;
            while (x != 0)
            {
                int b = System.Numerics.BitOperations.TrailingZeroCount(x);
                if (++count[b] > maxCount)
                    maxCount = count[b];
                x &= x - 1;
            }
        }

        long ans = 0;
        int limit = Math.Min(k, maxCount);
        for (int j = 0; j < limit; j++)
        {
            long value = 0;
            for (int b = 0; b < BITS; b++)
                if (count[b] > j)
                    value |= 1L << b;
            ans = (ans + value * value) % MOD;
        }
        return (int)ans;
    }
}
Was this solution helpful?

Related Problems