LeetCode/kt/prime-subtraction-operation.kt

class Solution {
    private fun getPrimes(nums: IntArray): BooleanArray {
        val primes = BooleanArray(1 + nums.max()) { true }
        primes[1] = false

        (2..sqrt(primes.size.toDouble()).toInt())
            .asSequence()
            .filter { primes[it] }
            .forEach { i ->
                (i * i..<primes.size step i).forEach { j ->
                    primes[j] = false
                }
            }

        return primes
    }

    fun primeSubOperation(nums: IntArray): Boolean {
        val primes = getPrimes(nums)

        var (current, i) = 1 to 0
        while (i < nums.size) {
            val difference = nums[i] - current
            if (difference < 0) {
                return false
            }

            if (primes[difference] || difference == 0) {
                i += 1
            }
            current += 1
        }

        return true
    }
}