aboutsummaryrefslogtreecommitdiff
path: root/apply_operations_arr_max_sqr.py
blob: e3b5adef7d13a299a7cd2556f4f0d3a324f01cab (plain)
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
# 2897. Apply Operations on Array to Maximize Sum of Squares
from collections import Counter

"""
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] & nums[j])' and 'nums[j]' to '(nums[i] | nums[j])'.
you have to choose 'k' elements from the final array and calcluate the sum of
their squares. return the maximum sum of squares you can achieve. since the
answer can be very large, return it modulo 10^9+7.
"""


class Solution(object):
    def maxSum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        cnt = Counter(i for j in nums for i in range(32) if j & (1 << i))
        return sum(sum((cnt[i] > j) << i for i in range(32)) ** 2 for j in range(k)) % (
            10**9 + 7
        )


if __name__ == "__main__":
    obj = Solution()
    print(obj.maxSum(nums=[2, 6, 5, 8], k=2))
    print(obj.maxSum(nums=[4, 5, 4, 7], k=3))