LeetCode Problem 3618

Difficulty: 🟡 Medium
Topics: Array, Math, Number Theory

Solution

LeetCode 3618 - Split Array by Prime Indices

Problem Understanding

The problem gives us an integer array nums and asks us to divide its elements into two groups based on their indices.

Array A contains all elements whose indices are prime numbers. Array B contains all remaining elements, meaning indices that are not prime. Remember that array indices are zero-based, so the first element has index 0, the second element has index 1, and so on.

After splitting the array, we compute:

• sum(A), the sum of all elements placed into A
• sum(B), the sum of all elements placed into B

The final answer is the absolute difference between these two sums:

$$|sum(A) - sum(B)|$$

The array itself does not need to be physically split. We only need the two sums.

The constraints are important:

• 1 <= nums.length <= 10^5
• -10^9 <= nums[i] <= 10^9

Since the array can contain up to one hundred thousand elements, we need an efficient way to determine whether each index is prime. A naive primality test for every index would still work in practice for this size, but there is a more efficient and cleaner solution using the Sieve of Eratosthenes.

The values themselves can be very large and can also be negative. Therefore, the running sums may become much larger than a 32-bit integer range. Using Python integers is safe because they are arbitrary precision. In Go, we should use int64 for the sums.

Some important edge cases include:

• Arrays of length 1, where there are no prime indices.
• Arrays of length 2, where indices 0 and 1 are both non-prime.
• Arrays containing negative numbers.
• Arrays where all values are zero.
• Large arrays where efficient prime detection matters.

Approaches

Brute Force

A straightforward approach is to iterate through every index and determine whether that index is prime by trial division.

For each index i, we test divisibility by every integer from 2 up to √i. If no divisor is found, the index is prime and its value contributes to sum(A). Otherwise, it contributes to sum(B).

This approach is correct because it directly follows the definition of a prime number. However, primality testing is repeated independently for every index. With up to 10^5 indices, this introduces unnecessary repeated work.

Key Insight

The indices range only from 0 to n - 1, where n <= 10^5.

Instead of testing each index separately, we can precompute all prime indices in one pass using the Sieve of Eratosthenes.

The sieve creates a boolean array where is_prime[i] tells us whether index i is prime. After building the sieve, we make a single pass through nums and add each value to the appropriate sum.

The Sieve of Eratosthenes computes primality for all numbers up to n - 1 in near-linear time, making it an ideal fit for the constraint size.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n√n)	O(1)	Test every index individually for primality
Optimal	O(n log log n)	O(n)	Use Sieve of Eratosthenes to precompute prime indices

Algorithm Walkthrough

1. Let n be the length of nums.
2. Create a boolean array is_prime of size n, initially marking all entries as prime.
3. If n > 0, mark index 0 as non-prime. If n > 1, mark index 1 as non-prime because neither 0 nor 1 is a prime number.
4. Run the Sieve of Eratosthenes. For every integer p from 2 up to √(n-1):

• If p is marked prime, then all multiples of p starting from p * p are composite.
• Mark those multiples as non-prime.

5. Initialize two running sums:

• prime_sum = 0
• non_prime_sum = 0

6. Traverse the array once.

• If is_prime[i] is true, add nums[i] to prime_sum.
• Otherwise, add nums[i] to non_prime_sum.

7. Return:

$$|prime_sum - non_prime_sum|$$

Why it works

The sieve guarantees that is_prime[i] is true exactly when index i is a prime number. Every array element is processed exactly once and placed into the correct conceptual group based on its index. Therefore, prime_sum equals the sum of all elements at prime indices and non_prime_sum equals the sum of all elements at non-prime indices. Taking the absolute difference between these two sums produces exactly the quantity required by the problem.

Python Solution

from typing import List

class Solution:
    def splitArray(self, nums: List[int]) -> int:
        n = len(nums)

        is_prime = [True] * n

        if n > 0:
            is_prime[0] = False
        if n > 1:
            is_prime[1] = False

        p = 2
        while p * p < n:
            if is_prime[p]:
                multiple = p * p
                while multiple < n:
                    is_prime[multiple] = False
                    multiple += p
            p += 1

        prime_sum = 0
        non_prime_sum = 0

        for i, value in enumerate(nums):
            if is_prime[i]:
                prime_sum += value
            else:
                non_prime_sum += value

        return abs(prime_sum - non_prime_sum)

The implementation begins by constructing the sieve for all possible indices. The array is_prime stores whether each index should belong to array A.

The sieve marks 0 and 1 as non-prime and then removes all composite numbers by marking multiples of each discovered prime.

Once the sieve is complete, the code performs a single traversal of nums. Depending on whether the current index is prime, the value contributes either to prime_sum or non_prime_sum.

Finally, the absolute difference between the two sums is returned.

Go Solution

func splitArray(nums []int) int64 {
	n := len(nums)

	isPrime := make([]bool, n)
	for i := range isPrime {
		isPrime[i] = true
	}

	if n > 0 {
		isPrime[0] = false
	}
	if n > 1 {
		isPrime[1] = false
	}

	for p := 2; p*p < n; p++ {
		if isPrime[p] {
			for multiple := p * p; multiple < n; multiple += p {
				isPrime[multiple] = false
			}
		}
	}

	var primeSum int64
	var nonPrimeSum int64

	for i, value := range nums {
		if isPrime[i] {
			primeSum += int64(value)
		} else {
			nonPrimeSum += int64(value)
		}
	}

	diff := primeSum - nonPrimeSum
	if diff < 0 {
		diff = -diff
	}

	return diff
}

The Go implementation follows the same algorithm. The main difference is that the running sums are stored as int64 to safely handle values that may exceed the range of a 32-bit integer. The final absolute value is computed manually because Go does not provide a built-in absolute value function for int64.

Worked Examples

Example 1

Input:

nums = [2, 3, 4]

Prime indices less than 3:

0 -> not prime
1 -> not prime
2 -> prime

Index	Value	Prime Index?	prime_sum	non_prime_sum
0	2	No	0	2
1	3	No	0	5
2	4	Yes	4	5

Final calculation:

$$|4 - 5| = 1$$

Answer:

1

Example 2

Input:

nums = [-1, 5, 7, 0]

Prime indices:

2, 3

Index	Value	Prime Index?	prime_sum	non_prime_sum
0	-1	No	0	-1
1	5	No	0	4
2	7	Yes	7	4
3	0	Yes	7	4

Final calculation:

$$|7 - 4| = 3$$

Answer:

3

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log log n)	Sieve construction dominates the running time
Space	O(n)	Boolean sieve stores primality information for every index

The Sieve of Eratosthenes runs in O(n log log n) time, which is effectively linear for the given constraint size. After building the sieve, we perform one additional linear scan through the array. The extra memory comes from the boolean array used to store primality information.

Test Cases

sol = Solution()

assert sol.splitArray([2, 3, 4]) == 1  # example 1
assert sol.splitArray([-1, 5, 7, 0]) == 3  # example 2

assert sol.splitArray([5]) == 5  # length 1, no prime indices
assert sol.splitArray([10, 20]) == 30  # length 2, no prime indices

assert sol.splitArray([1, 2, 3]) == 0  # balanced sums

assert sol.splitArray([0, 0, 0, 0, 0]) == 0  # all zeros

assert sol.splitArray([-5, -10, -20]) == 5  # negative values

assert sol.splitArray([1, 1, 1, 1, 1]) == 1  # prime indices 2 and 3

assert sol.splitArray([1000000000, -1000000000, 1000000000]) == 1000000000  # large magnitudes

assert sol.splitArray([1, 2, 3, 4, 5, 6, 7, 8]) == 10  # multiple prime indices

Test Summary

Test	Why
[2,3,4]	First provided example
[-1,5,7,0]	Second provided example
[5]	Minimum array length
[10,20]	No prime indices exist
[1,2,3]	Difference becomes zero
[0,0,0,0,0]	All sums remain zero
[-5,-10,-20]	Negative number handling
[1,1,1,1,1]	Multiple prime indices with equal values
[1000000000,-1000000000,1000000000]	Large integer magnitudes
[1,2,3,4,5,6,7,8]	General mixed case

Edge Cases

One important edge case is an array of length one. The only index is 0, which is not prime. Therefore, every element belongs to the non-prime group. A buggy implementation might accidentally treat 0 as prime or fail when constructing the sieve. The solution explicitly marks index 0 as non-prime.

Another important edge case is an array of length two. The indices are 0 and 1, and neither is prime. Some implementations incorrectly assume that 1 is prime. The algorithm explicitly marks both 0 and 1 as non-prime before running the sieve.

Arrays containing negative numbers are another common source of mistakes. Since the problem asks for the absolute difference of sums, intermediate sums may become negative. The implementation accumulates the true signed sums and only applies the absolute value at the end, ensuring correct behavior regardless of the sign of individual elements.

A final edge case involves very large element values, up to 10^9. Since there can be up to 10^5 elements, the total sum may reach approximately 10^14 in magnitude. Python handles this automatically. In Go, the implementation uses int64 for the sums, preventing overflow that would occur with a 32-bit integer type.