2862. Maximum Element-Sum of a Complete Subset of Indices

Problem Link

Description

---

You are given a 1-indexed array nums. Your task is to select a complete subset from nums where every pair of selected indices multiplied is a perfect square,. i. e. if you select ai and aj, i * j must be a perfect square.

Return the sum of the complete subset with the maximum sum.

 

Example 1:

Input: nums = [8,7,3,5,7,2,4,9]

Output: 16

Explanation:

We select elements at indices 2 and 8 and 2 * 8 is a perfect square.

Example 2:

Input: nums = [8,10,3,8,1,13,7,9,4]

Output: 20

Explanation:

We select elements at indices 1, 4, and 9. 1 * 4, 1 * 9, 4 * 9 are perfect squares.

 

Constraints:

• 1 <= n == nums.length <= 104
• 1 <= nums[i] <= 109

Solution

---

Python3

class Solution:
   def maximumSum(self, nums: List[int]) -> int:
        N = len(nums)
        res = 0
        squares = []
        i = 1
        
        while i <= N:
            squares.append(i * i)
            i += 1
 
        # 1) perfect sqaure * perfect square = perfect square -> a^2 * b^2 = a^2b^2
        # 2) since i is constant, pair1 = i^1 * a^2, pair2 = i^1 * b^2
        # 3) their product is i^2 * (ab) ^ 2 which warrants a perfect square
 
        for i, x in enumerate(nums, start = 1):
            curr = 0
 
            for sq in squares:
                if i * sq - 1 >= N: break
                
                curr += nums[i * sq - 1]
 
            res = max(res, curr)
 
        return res