LeetCode in Kotlin

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

import kotlin.math.abs

class Solution {
    fun splitArray(nums: IntArray): Long {
        val n = nums.size
        val isPrime = sieve(n)
        var sumA: Long = 0
        var sumB: Long = 0
        for (i in 0..<n) {
            if (isPrime[i]) {
                sumA += nums[i].toLong()
            } else {
                sumB += nums[i].toLong()
            }
        }
        return abs(sumA - sumB)
    }

    // Sieve of Eratosthenes to find all prime indices up to n
    private fun sieve(n: Int): BooleanArray {
        val isPrime = BooleanArray(n)
        if (n > 2) {
            isPrime[2] = true
        }
        run {
            var i = 3
            while (i < n) {
                isPrime[i] = true
                i += 2
            }
        }
        if (n > 2) {
            isPrime[2] = true
        }
        var i = 3
        while (i * i < n) {
            if (isPrime[i]) {
                var j = i * i
                while (j < n) {
                    isPrime[j] = false
                    j += i * 2
                }
            }
            i += 2
        }
        isPrime[0] = false
        if (n > 1) {
            isPrime[1] = false
        }
        return isPrime
    }
}