LeetCode-in-Java

3618. Split Array by Prime Indices

Medium

You are given an integer array nums.

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

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:

Example 2:

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

Output: 3

Explanation:

Constraints:

Solution

public class Solution {
    public long splitArray(int[] nums) {
        int n = nums.length;
        boolean[] isPrime = sieve(n);
        long sumA = 0;
        long sumB = 0;
        for (int i = 0; i < n; i++) {
            if (isPrime[i]) {
                sumA += nums[i];
            } else {
                sumB += nums[i];
            }
        }
        return Math.abs(sumA - sumB);
    }

    // Sieve of Eratosthenes to find all prime indices up to n
    private boolean[] sieve(int n) {
        boolean[] isPrime = new boolean[n];
        if (n > 2) {
            isPrime[2] = true;
        }
        for (int i = 3; i < n; i += 2) {
            isPrime[i] = true;
        }
        if (n > 2) {
            isPrime[2] = true;
        }
        for (int i = 3; i * i < n; i += 2) {
            if (isPrime[i]) {
                for (int j = i * i; j < n; j += i * 2) {
                    isPrime[j] = false;
                }
            }
        }
        isPrime[0] = false;
        if (n > 1) {
            isPrime[1] = false;
        }
        return isPrime;
    }
}