Skip to main content
Talentd

Split Array by Prime Indices - Solution & Explanation

MediumArrayMathNumber Theory7 min read
Practice this problem

Problem Statement

You are given an integer array nums.

Split nums into two arrays A and B using the following rule:

  • Elements at prime indices in nums must go into array A.
  • All other elements must go into array B.

Return the absolute difference between the sums of the two arrays: |sum(A) - sum(B)|.

Note: An empty array has a sum of 0.

 

Example 1:

Input: nums = [2,3,4]

Output: 1

Explanation:

  • The only prime index in the array is 2, so nums[2] = 4 is placed in array A.
  • The remaining elements, nums[0] = 2 and nums[1] = 3 are placed in array B.
  • sum(A) = 4, sum(B) = 2 + 3 = 5.
  • The absolute difference is |4 - 5| = 1.

Example 2:

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

Output: 3

Explanation:

  • The prime indices in the array are 2 and 3, so nums[2] = 7 and nums[3] = 0 are placed in array A.
  • The remaining elements, nums[0] = -1 and nums[1] = 5 are placed in array B.
  • sum(A) = 7 + 0 = 7, sum(B) = -1 + 5 = 4.
  • The absolute difference is |7 - 4| = 3.

 

Constraints:

  • 1 <= nums.length <= 105
  • -109 <= nums[i] <= 109

Approach Overview

Problem Overview: You are given an array and need to split or process its elements based on whether their indices are prime numbers. The core task is identifying prime indices efficiently and then simulating the required split or aggregation using those indices.

Approach 1: Naive Prime Check + Simulation (O(n * sqrt(n)) time, O(1) space)

The straightforward approach checks whether each index is prime by testing divisibility from 2 to sqrt(i). For every array index, run the prime check and then place or process the element accordingly. This method works for small inputs but becomes inefficient because each index requires a separate primality test. The algorithm repeatedly performs expensive checks even though many results could be reused.

Approach 2: Sieve of Eratosthenes + Simulation (O(n log log n) time, O(n) space)

The efficient solution precomputes all prime indices up to n - 1 using the Sieve of Eratosthenes. Create a boolean array where each position indicates whether the index is prime. Start by marking multiples of each prime number as non‑prime, which builds the complete prime table in O(n log log n) time. After building the sieve, iterate through the array once and check the precomputed prime flag for each index.

If the index is prime, process the element in the "prime index" group; otherwise place it in the "non‑prime index" group. The simulation step is linear and consists of simple array iteration and conditional checks. Because primality is already computed, each lookup is O(1), making the total runtime dominated by the sieve preprocessing.

This pattern frequently appears in problems combining array traversal with math utilities. Precomputing primes with a sieve is a standard optimization in number theory tasks where repeated prime checks are required.

Recommended for interviews: The Sieve of Eratosthenes + simulation approach is the expected solution. Interviewers want to see that you avoid repeated primality checks and instead precompute primes efficiently. Mentioning the naive method shows you understand the baseline, but using the sieve demonstrates stronger algorithmic optimization and familiarity with number‑theory preprocessing techniques.

Solution

We can use the Sieve of Eratosthenes to preprocess all prime numbers in the range [0, 10^5]. Then we iterate through the array nums. For nums[i], if i is a prime number, we add nums[i] to the answer; otherwise, we add -nums[i] to the answer. Finally, we return the absolute value of the answer.

Ignoring the preprocessing time and space, the time complexity is O(n), where n is the length of the array nums, and the space complexity is O(1).

Code

PythonJavaC++GoTypeScript

Python

Java

C++

Go

TypeScript

Try this approach in the editor →

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
Naive Prime Check + SimulationO(n * sqrt(n))O(1)When constraints are small or when demonstrating baseline logic before optimization
Sieve of Eratosthenes + SimulationO(n log log n)O(n)General case for large arrays where many prime index checks are required

Video Solution

Split Array by Prime Indices | Biweekly Contest 161 | Q1 | Leetcode 3618 • ExpertFunda • 270 views views

Watch 9 more video solutions →

Frequently Asked Questions

Is Split Array by Prime Indices easy or hard?
Split Array by Prime Indices is generally considered a medium‑difficulty problem. The array iteration itself is straightforward, but recognizing that repeated prime checks should be replaced with a sieve optimization is the key insight.
Split Array by Prime Indices Python/Java solution
The typical implementation first constructs a sieve array up to n−1, marking prime indices. Then a simple loop iterates through the array and processes elements depending on whether the index is prime. The same logic works across Python, Java, C++, Go, and TypeScript.
How to solve Split Array by Prime Indices in O(n)?
After precomputing prime indices using the sieve, the actual processing step runs in O(n). Iterate through the array once and check a boolean prime table to decide whether the current index belongs to the prime group or the non‑prime group. Each lookup is constant time.
What is the best approach for Split Array by Prime Indices?
The most efficient approach uses the Sieve of Eratosthenes to precompute all prime indices up to the array length. After building the sieve in O(n log log n) time, iterate through the array once and process elements based on whether their index is marked prime. This avoids repeated primality checks and keeps the simulation step linear.
Is Split Array by Prime Indices asked at Google/Amazon/Meta?
Problems involving prime preprocessing and array traversal appear frequently in interviews at companies like Google, Amazon, and Meta. Variations often require identifying prime indices, using sieve techniques, or combining number theory with efficient array processing.
What data structure is used in Split Array by Prime Indices?
The core structure is a boolean array (or vector) used by the Sieve of Eratosthenes to mark prime numbers. The input array is then traversed while referencing this table to determine whether each index is prime.
What is the time complexity of Split Array by Prime Indices?
Using the optimal approach, the time complexity is O(n log log n) due to the Sieve of Eratosthenes preprocessing. The final array traversal is O(n). Space complexity is O(n) for storing the prime flags.

Ready to solve this problem?

Practice Split Array by Prime Indices with our built-in code editor and test cases.

Practice on FleetCode
Previous

Two Sum

Next

Median of Two Sorted Arrays