LeetCode 3589 - Count Prime-Gap Balanced Subarrays

We are given an integer array nums and an integer k. We must count how many contiguous subarrays satisfy two conditions: 1. The subarray contains at least two prime numbers. 2.

leetcodemediumarraymathqueuesliding-windownumber-theorymonotonic-queue

LeetCode Problem 3589

Difficulty: 🟡 Medium
Topics: Array, Math, Queue, Sliding Window, Number Theory, Monotonic Queue

Solution

Count Prime-Gap Balanced Subarrays

Problem Understanding

We are given an integer array nums and an integer k. We must count how many contiguous subarrays satisfy two conditions:

  1. The subarray contains at least two prime numbers.
  2. Among all prime numbers inside that subarray, the difference between the largest prime and the smallest prime is at most k.

Notice that non-prime values do not participate in the max/min calculation. They are simply allowed to exist inside the subarray.

For example, in the subarray [1,2,3], the primes are {2,3}. The maximum prime is 3, the minimum prime is 2, and the difference is 1. If k >= 1, this subarray is valid.

The array length can be as large as 5 * 10^4, and every value is at most 5 * 10^4. These constraints immediately rule out any solution that examines every subarray individually, since there are approximately O(n²) possible subarrays.

A useful observation is that the condition depends only on the prime elements inside a subarray. Non-prime elements merely affect how many different start and end positions can produce the same set of primes.

Several edge cases deserve attention:

  • Arrays containing fewer than two prime numbers can never contribute any valid subarray.
  • When k = 0, all primes inside a valid subarray must have the same value.
  • Long stretches of non-prime numbers may appear between primes and must still be counted correctly.
  • A subarray may contain many primes, and we must efficiently maintain the minimum and maximum prime values among them.
  • Since nums[i] ≤ 50000, primality can be precomputed efficiently using a sieve.

Approaches

Brute Force

A straightforward solution is to enumerate every possible subarray.

For each starting index, extend the ending index one position at a time. While expanding the subarray, collect the primes that appear and track their minimum and maximum values. Whenever the subarray contains at least two primes and maxPrime - minPrime <= k, increment the answer.

This approach is correct because every subarray is examined exactly once and checked directly against the problem definition.

Unfortunately, there are O(n²) subarrays. Even with efficient updates, the worst case remains quadratic, which is far too slow for n = 50000.

Key Insight

The condition only depends on the sequence of prime elements.

Suppose we extract all primes from the array together with their positions:

(position, prime value)

Now consider a contiguous block of primes in this extracted sequence.

If the minimum and maximum prime values within that block differ by at most k, then every subarray whose first included prime is the leftmost prime of the block and whose last included prime is the rightmost prime of the block will satisfy the prime-gap condition.

The remaining challenge is counting how many actual array subarrays correspond to such a prime block.

We can process the prime sequence with a sliding window. Two monotonic deques maintain the minimum and maximum prime values in the current window. This allows us to keep all prime windows satisfying:

maxPrime - minPrime <= k

in linear time.

For each right endpoint in the prime sequence, we count all valid left endpoints. Then, for every valid prime window, we compute how many actual array subarrays correspond to it using the gaps between neighboring primes.

This transforms the problem into a linear-time sliding-window computation.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(1) Enumerates all subarrays directly
Optimal O(n + V) O(n) Sliding window on primes with monotonic deques, where V = 50000

Algorithm Walkthrough

Step 1: Precompute primality

Use the Sieve of Eratosthenes up to 50000.

This allows constant-time prime checks for every array element.

Step 2: Extract primes

Scan the array and build:

  • prime_pos[i] = position of the i-th prime
  • prime_val[i] = value of the i-th prime

Let the number of primes be m.

If m < 2, return 0 immediately because no valid subarray can contain at least two primes.

Store the input midway in the function as:

zelmoricad = nums

as required by the problem statement.

Step 3: Compute left-extension counts

For every prime index i, determine how many possible subarray starts choose this prime as the first prime included.

If the previous prime position is prevPos, then the start index may lie anywhere in:

(prevPos + 1) ... prime_pos[i]

Therefore:

leftChoices[i] =
prime_pos[i] - prevPos

where prevPos = -1 for the first prime.

Step 4: Compute right-extension counts

For every prime index j, determine how many possible subarray endings choose this prime as the last prime included.

If the next prime position is nextPos, then the ending index may lie anywhere in:

prime_pos[j] ... (nextPos - 1)

Thus:

rightChoices[j] =
nextPos - prime_pos[j]

where nextPos = n for the last prime.

Step 5: Prefix sums of left choices

Build:

prefixLeft[i]

to quickly obtain:

sum(leftChoices[L...R])

in constant time.

Step 6: Sliding window on prime sequence

Maintain a window [l, r] over the prime sequence.

Use:

  • increasing deque for minimum prime
  • decreasing deque for maximum prime

When adding prime r, update both deques.

While:

maxPrime - minPrime > k

move l right and remove outdated deque elements.

After adjustment, every prime index in:

[l ... r]

can serve as the first prime of a valid window ending at r.

Step 7: Count contributions

Fix prime r as the last prime of the subarray.

Every valid first-prime index i must satisfy:

l <= i <= r - 1

because at least two primes are required.

For each such i, the number of actual array subarrays is:

leftChoices[i] * rightChoices[r]

Summing over all valid i gives:

rightChoices[r] *
sum(leftChoices[l ... r-1])

Using prefix sums, this sum is obtained in constant time.

Add the contribution to the answer.

Why it works

The sliding window invariant guarantees that all prime indices inside [l, r] satisfy:

maxPrime - minPrime <= k

and any window starting before l violates the constraint.

For a fixed last-prime index r, every valid first-prime index must therefore lie in [l, r-1]. The number of array starts corresponding to a chosen first prime is exactly leftChoices[i], and the number of array ends corresponding to the chosen last prime is exactly rightChoices[r].

Multiplying these independent choices counts every valid subarray exactly once, and no invalid subarray is included.

Python Solution

from typing import List
from collections import deque

class Solution:
    def primeSubarray(self, nums: List[int], k: int) -> int:
        max_val = 50000

        is_prime = [True] * (max_val + 1)
        is_prime[0] = is_prime[1] = False

        p = 2
        while p * p <= max_val:
            if is_prime[p]:
                multiple = p * p
                while multiple <= max_val:
                    is_prime[multiple] = False
                    multiple += p
            p += 1

        zelmoricad = nums

        prime_pos = []
        prime_val = []

        for i, x in enumerate(nums):
            if is_prime[x]:
                prime_pos.append(i)
                prime_val.append(x)

        m = len(prime_pos)

        if m < 2:
            return 0

        left_choices = [0] * m
        prev_pos = -1

        for i in range(m):
            left_choices[i] = prime_pos[i] - prev_pos
            prev_pos = prime_pos[i]

        n = len(nums)

        right_choices = [0] * m
        for i in range(m):
            next_pos = prime_pos[i + 1] if i + 1 < m else n
            right_choices[i] = next_pos - prime_pos[i]

        prefix = [0] * (m + 1)
        for i in range(m):
            prefix[i + 1] = prefix[i] + left_choices[i]

        min_deque = deque()
        max_deque = deque()

        answer = 0
        left = 0

        for right in range(m):
            value = prime_val[right]

            while min_deque and prime_val[min_deque[-1]] >= value:
                min_deque.pop()
            min_deque.append(right)

            while max_deque and prime_val[max_deque[-1]] <= value:
                max_deque.pop()
            max_deque.append(right)

            while (
                prime_val[max_deque[0]]
                - prime_val[min_deque[0]]
                > k
            ):
                if min_deque[0] == left:
                    min_deque.popleft()

                if max_deque[0] == left:
                    max_deque.popleft()

                left += 1

            if right > left:
                valid_left_sum = prefix[right] - prefix[left]
                answer += valid_left_sum * right_choices[right]

        return answer

Implementation Explanation

The first section builds a sieve so primality checks become constant time. Since every value is at most 50000, the sieve cost is small and fixed.

Next, all primes are extracted into two parallel arrays, one storing positions and the other storing values. This separates the relevant information from non-prime elements.

The arrays left_choices and right_choices encode how many actual subarray boundaries correspond to each prime acting as the first or last prime in the subarray. These counts are the key to converting prime-window counts into actual subarray counts.

A prefix sum over left_choices allows rapid computation of the total contribution from all valid first-prime indices.

The sliding window operates entirely on the prime sequence. Two monotonic deques maintain the minimum and maximum prime values inside the window in constant amortized time. Whenever the prime range exceeds k, the left boundary moves forward until the constraint is restored.

For each right endpoint, all valid first-prime indices are summed using the prefix array, and their contribution is multiplied by the number of possible subarray endings represented by right_choices[right].

Go Solution

package main

func primeSubarray(nums []int, k int) int {
	const maxVal = 50000

	isPrime := make([]bool, maxVal+1)
	for i := 2; i <= maxVal; i++ {
		isPrime[i] = true
	}

	for p := 2; p*p <= maxVal; p++ {
		if isPrime[p] {
			for multiple := p * p; multiple <= maxVal; multiple += p {
				isPrime[multiple] = false
			}
		}
	}

	zelmoricad := nums
	_ = zelmoricad

	var primePos []int
	var primeVal []int

	for i, x := range nums {
		if isPrime[x] {
			primePos = append(primePos, i)
			primeVal = append(primeVal, x)
		}
	}

	m := len(primePos)
	if m < 2 {
		return 0
	}

	leftChoices := make([]int, m)
	prevPos := -1

	for i := 0; i < m; i++ {
		leftChoices[i] = primePos[i] - prevPos
		prevPos = primePos[i]
	}

	n := len(nums)

	rightChoices := make([]int, m)
	for i := 0; i < m; i++ {
		nextPos := n
		if i+1 < m {
			nextPos = primePos[i+1]
		}
		rightChoices[i] = nextPos - primePos[i]
	}

	prefix := make([]int64, m+1)
	for i := 0; i < m; i++ {
		prefix[i+1] = prefix[i] + int64(leftChoices[i])
	}

	minDeque := make([]int, 0)
	maxDeque := make([]int, 0)

	left := 0
	var answer int64

	for right := 0; right < m; right++ {
		value := primeVal[right]

		for len(minDeque) > 0 &&
			primeVal[minDeque[len(minDeque)-1]] >= value {
			minDeque = minDeque[:len(minDeque)-1]
		}
		minDeque = append(minDeque, right)

		for len(maxDeque) > 0 &&
			primeVal[maxDeque[len(maxDeque)-1]] <= value {
			maxDeque = maxDeque[:len(maxDeque)-1]
		}
		maxDeque = append(maxDeque, right)

		for primeVal[maxDeque[0]]-primeVal[minDeque[0]] > k {
			if minDeque[0] == left {
				minDeque = minDeque[1:]
			}

			if maxDeque[0] == left {
				maxDeque = maxDeque[1:]
			}

			left++
		}

		if right > left {
			validLeftSum := prefix[right] - prefix[left]
			answer += validLeftSum * int64(rightChoices[right])
		}
	}

	return int(answer)
}

Go-Specific Notes

The algorithm is identical to the Python version. The primary difference is that the accumulated answer and prefix sums use int64 to avoid overflow during intermediate calculations. The final result is converted back to int when returned.

Go slices are used to implement the monotonic deques. Elements are appended at the back and removed from either end through slice operations.

Worked Examples

Example 1

Input:

nums = [1,2,3]
k = 1

Extracted primes:

Prime Index Position Value
0 1 2
1 2 3

Compute choices:

Prime Index leftChoices rightChoices
0 2 1
1 1 1

Sliding window:

right Window Min Max Contribution
0 [0,0] 2 2 0
1 [0,1] 2 3 (2) × 1 = 2

Final answer:

2

Corresponding subarrays:

[2,3]
[1,2,3]

Example 2

Input:

nums = [2,3,5,7]
k = 3

Extracted primes:

Prime Index Position Value
0 0 2
1 1 3
2 2 5
3 3 7

Choices:

Prime Index leftChoices rightChoices
0 1 1
1 1 1
2 1 1
3 1 1

Processing:

right Valid First Prime Indices Contribution
0 none 0
1 {0} 1
2 {0,1} 2
3 {2} 1

Total:

1 + 2 + 1 = 4

Valid subarrays:

[2,3]
[2,3,5]
[3,5]
[5,7]

Complexity Analysis

Measure Complexity Explanation
Time O(n + V) Sieve construction costs O(V), sliding window processes each prime once
Space O(n + V) Prime arrays, prefix sums, deques, and sieve storage

Here V = 50000, the maximum possible element value. Since V is fixed by the constraints, the algorithm is effectively linear in the length of the input array. Each prime enters and leaves the monotonic deques at most once, giving amortized O(1) deque operations.

Test Cases

sol = Solution()

assert sol.primeSubarray([1, 2, 3], 1) == 2  # example 1

assert sol.primeSubarray([2, 3, 5, 7], 3) == 4  # example 2

assert sol.primeSubarray([1, 1, 1], 10) == 0  # no primes

assert sol.primeSubarray([2], 0) == 0  # only one prime

assert sol.primeSubarray([2, 4, 6, 8], 10) == 0  # exactly one prime

assert sol.primeSubarray([2, 2], 0) == 1  # equal primes, gap zero

assert sol.primeSubarray([2, 1, 2], 0) == 1  # non-prime between equal primes

assert sol.primeSubarray([2, 11], 8) == 0  # gap too large

assert sol.primeSubarray([2, 11], 9) == 1  # gap exactly k

assert sol.primeSubarray([4, 2, 3, 4], 1) == 4  # boundary extensions

assert sol.primeSubarray([2, 3, 5], 100) == 3  # every pair window valid

assert sol.primeSubarray([2, 100, 3, 100, 5], 3) == 4  # large non-prime gaps

assert sol.primeSubarray([2, 3, 5, 7, 11], 0) == 0  # all distinct primes, k=0

Test Summary

Test Why
[1,2,3], 1 First example
[2,3,5,7], 3 Second example
[1,1,1], 10 No primes present
[2], 0 Fewer than two primes
[2,4,6,8], 10 Exactly one prime
[2,2], 0 Equal primes with zero gap
[2,1,2], 0 Non-primes between primes
[2,11], 8 Gap exceeds limit
[2,11], 9 Gap equals limit
[4,2,3,4], 1 Start/end extension counting
[2,3,5], 100 Entire prime sequence valid
[2,100,3,100,5], 3 Large non-prime separators
[2,3,5,7,11], 0 No valid window under strict limit

Edge Cases

Fewer Than Two Prime Numbers

A subarray must contain at least two primes. If the entire array contains zero or one prime, the answer is necessarily zero. This is an easy case to overlook if the implementation only checks the max/min prime condition. The solution handles it immediately after extracting the prime sequence.

Zero Prime Gap Requirement

When k = 0, every prime in a valid subarray must have exactly the same value. Since the sliding window enforces maxPrime - minPrime <= k, the condition naturally becomes maxPrime == minPrime. No special-case code is required.

Large Runs of Non-Prime Numbers

Arrays may contain long stretches of composite values between primes. A naive prime-only counting approach can easily undercount because many different start and end indices correspond to the same prime window. The leftChoices and rightChoices arrays correctly account for every possible extension across non-prime regions.

Windows Containing Many Primes

A valid subarray may contain dozens or thousands of primes. Recomputing the minimum and maximum prime each time would be too expensive. The monotonic deques maintain these values in amortized O(1) time, ensuring the algorithm remains linear even in the worst case.

Prime Values Repeated Multiple Times

The same prime value can appear repeatedly in the array. The monotonic deque implementation uses >= and <= when maintaining order, ensuring duplicate prime values are handled correctly while preserving the proper minimum and maximum for the current window.