Skip to content

3618. Split Array by Prime Indices

Description

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

Solutions

Solution 1: Sieve of Eratosthenes + Simulation

Thinking

The answer is the absolute difference between the sum at prime indices and the sum at composite indices, i.e. a signed sum over \(i\). Indices reach \(10^5\), so trial division per index is wasteful.

The sieve of Eratosthenes marks primality on \([0,10^5]\). Walk the array, add \(x\) at a prime index and \(-x\) otherwise, then take the absolute value.

The sieve can be reused; each query is a single linear scan.

We can use the Sieve of Eratosthenes to preprocess all prime numbers in the range \([0, 10^5]\). Then we iterate through the array \(\textit{nums}\). For \(\textit{nums}[i]\), if \(i\) is a prime number, we add \(\textit{nums}[i]\) to the answer; otherwise, we add \(-\textit{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 \(\textit{nums}\), and the space complexity is \(O(1)\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
m = 10**5 + 10
primes = [True] * m
primes[0] = primes[1] = False
for i in range(2, m):
    if primes[i]:
        for j in range(i + i, m, i):
            primes[j] = False


class Solution:
    def splitArray(self, nums: List[int]) -> int:
        return abs(sum(x if primes[i] else -x for i, x in enumerate(nums)))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution {
    private static final int M = 100000 + 10;
    private static boolean[] primes = new boolean[M];

    static {
        for (int i = 0; i < M; i++) {
            primes[i] = true;
        }
        primes[0] = primes[1] = false;

        for (int i = 2; i < M; i++) {
            if (primes[i]) {
                for (int j = i + i; j < M; j += i) {
                    primes[j] = false;
                }
            }
        }
    }

    public long splitArray(int[] nums) {
        long ans = 0;
        for (int i = 0; i < nums.length; ++i) {
            ans += primes[i] ? nums[i] : -nums[i];
        }
        return Math.abs(ans);
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
const int M = 1e5 + 10;
bool primes[M];
auto init = [] {
    memset(primes, true, sizeof(primes));
    primes[0] = primes[1] = false;
    for (int i = 2; i < M; ++i) {
        if (primes[i]) {
            for (int j = i + i; j < M; j += i) {
                primes[j] = false;
            }
        }
    }
    return 0;
}();

class Solution {
public:
    long long splitArray(vector<int>& nums) {
        long long ans = 0;
        for (int i = 0; i < nums.size(); ++i) {
            ans += primes[i] ? nums[i] : -nums[i];
        }
        return abs(ans);
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
const M = 100000 + 10

var primes [M]bool

func init() {
    for i := 0; i < M; i++ {
        primes[i] = true
    }
    primes[0], primes[1] = false, false

    for i := 2; i < M; i++ {
        if primes[i] {
            for j := i + i; j < M; j += i {
                primes[j] = false
            }
        }
    }
}

func splitArray(nums []int) (ans int64) {
    for i, num := range nums {
        if primes[i] {
            ans += int64(num)
        } else {
            ans -= int64(num)
        }
    }
    return max(ans, -ans)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
const M = 100000 + 10;
const primes: boolean[] = Array(M).fill(true);

const init = (() => {
    primes[0] = primes[1] = false;

    for (let i = 2; i < M; i++) {
        if (primes[i]) {
            for (let j = i + i; j < M; j += i) {
                primes[j] = false;
            }
        }
    }
})();

function splitArray(nums: number[]): number {
    let ans = 0;
    for (let i = 0; i < nums.length; i++) {
        ans += primes[i] ? nums[i] : -nums[i];
    }
    return Math.abs(ans);
}

Comments