Welcome to Subscribe On Youtube

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

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)$.

  • 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);
    }
}


  • 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);
    }
};

  • 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)))


  • 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)
}


  • 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);
}

All Problems

All Solutions