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