Maximum Element-Sum of a Complete Subset of Indices

HARD

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

Approaches

Checkout 2 different approaches to solve Maximum Element-Sum of a Complete Subset of Indices. Click on different approaches to view the approach and algorithm in detail.

Optimized Grouping using a Sieve

This is an optimized version of the first approach. It recognizes that calculating the square-free part for each index from scratch is inefficient. By precomputing the square-free parts for all numbers from 1 to n using a sieve, we can reduce the time complexity significantly. The sieve efficiently removes all perfect square factors from each number in nearly linear time. Once the precomputation is done, the problem is solved by grouping and summing the elements, just like in the previous approach, but with a much faster lookup for the square-free parts.

Algorithm

  • This approach uses the same fundamental principle of grouping indices by their square-free part.
  • To improve efficiency, it avoids re-calculating the square-free part for each index individually.
  • Instead, it precomputes the square-free parts for all numbers from 1 to n using a sieve-like algorithm.
  • The sieve works by iterating through numbers j and for each j, removing the perfect square factor j*j from all its multiples up to n.
  • After this O(n) precomputation, the main logic proceeds as before: iterate through indices, look up the precomputed square-free part, and aggregate sums in a hash map.

The bottleneck in the previous approach is the repeated calculation within the main loop. We can optimize this by pre-calculating all necessary square-free parts at once.

Algorithm Steps:

  1. Precomputation with a Sieve:
    • Create an integer array sf of size n + 1.
    • Initialize sf[i] = i for all i from 1 to n.
    • Iterate with a number j from 2 up to n. If j*j > n, we can stop.
    • For each j, let square = j * j.
    • Iterate through all multiples of square up to n (i.e., m = square, 2*square, 3*square, ...).
    • For each multiple m, repeatedly divide sf[m] by square until it's no longer divisible. This effectively removes the perfect square factor j*j from sf[m].
  2. Grouping and Summation:
    • Initialize a HashMap<Integer, Long> called groupSums.
    • Iterate with an index i from 1 to n.
    • The square-free part of i is now a simple lookup: core = sf[i].
    • Add nums[i-1] to the sum for the group core in the groupSums map.
  3. Find Maximum:
    • Find and return the maximum value present in the groupSums map.

This precomputation step changes the overall time complexity from O(n * sqrt(n)) to nearly O(n).

import java.util.HashMap;
import java.util.Map;

class Solution {
    public long maximumSum(int[] nums) {
        int n = nums.length;
        
        // Step 1: Precompute square-free parts for 1 to n using a sieve
        int[] sf = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            sf[i] = i;
        }
        
        for (int j = 2; j * j <= n; j++) {
            int square = j * j;
            for (int m = square; m <= n; m += square) {
                while (sf[m] % square == 0) {
                    sf[m] /= square;
                }
            }
        }
        
        // Step 2: Group sums based on the precomputed square-free parts
        Map<Integer, Long> groupSums = new HashMap<>();
        for (int i = 1; i <= n; i++) {
            int core = sf[i]; // O(1) lookup
            long currentVal = nums[i - 1];
            groupSums.put(core, groupSums.getOrDefault(core, 0L) + currentVal);
        }
        
        // Step 3: Find the maximum sum among all groups
        long maxSum = 0;
        for (long sum : groupSums.values()) {
            maxSum = Math.max(maxSum, sum);
        }
        
        return maxSum;
    }
}

Complexity Analysis

Time Complexity: O(n log log n) - The sieve precomputation step runs in `O(n log log n)` or `O(n)` time, which is much faster than `O(n * sqrt(n))`. The summation part runs in `O(n)`. Therefore, the total time complexity is dominated by the sieve and is effectively linear.Space Complexity: O(n) - This approach requires `O(n)` space for the `sf` array used in the sieve, plus `O(n)` space for the hash map in the worst case.

Pros and Cons

Pros:
  • Highly efficient with a nearly linear time complexity.
  • Optimal solution for the given constraints.
Cons:
  • Requires additional O(n) space for the sieve array.
  • The sieve logic, while efficient, can be slightly more complex to understand and implement compared to the direct calculation.

Code Solutions

Checking out 3 solutions in different languages for Maximum Element-Sum of a Complete Subset of Indices. Click on different languages to view the code.

class Solution {
public
  long maximumSum(List<Integer> nums) {
    long ans = 0;
    int n = nums.size();
    boolean[] used = new boolean[n + 1];
    int bound = (int)Math.floor(Math.sqrt(n));
    int[] squares = new int[bound + 1];
    for (int i = 1; i <= bound + 1; i++) {
      squares[i - 1] = i * i;
    }
    for (int i = 1; i <= n; i++) {
      long res = 0;
      int idx = 0;
      int curr = i * squares[idx];
      while (curr <= n) {
        res += nums.get(curr - 1);
        curr = i * squares[++idx];
      }
      ans = Math.max(ans, res);
    }
    return ans;
  }
}

Video Solution

Watch the video walkthrough for Maximum Element-Sum of a Complete Subset of Indices



Patterns:

MathNumber Theory

Data Structures:

Array

Subscribe to Scale Engineer newsletter

Learn about System Design, Software Engineering, and interview experiences every week.

No spam, unsubscribe at any time.