Leetcode 3589. Count Prime-Gap Balanced Subarrays

• Leetcode 3589. Count Prime-Gap Balanced Subarrays
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3589. Count Prime-Gap Balanced Subarrays

1. 解题思路

这一题思路上就是一个滑动窗口的思路。

我们首先找出所有的质数与其所在的位置，然后考察每一个质数作为选取的subarray的最后一个质数时其可以构成的subarray的数目。此时，我们只需要维护一个有序数组，其内容是截至当前质数情况下所有前方连续的差值不超过k的所有质数。

如此，我们就能够获得一个滑动窗口，其右边界的选择就是当前质数到其下一个质数之间的所有位置，而其左边界则是其前一个质数与滑动窗口左侧第一个质数之间的所有位置。

2. 代码实现

给出python代码实现如下：

def get_primes(n):
    status = [0 for _ in range(n+1)]
    primes = set()
    for i in range(2, n+1):
        if status[i] == 1:
            continue
        primes.add(i)
        for j in range(i, n+1, i):
            status[j] = 1
    return primes

PRIMES = get_primes(50000)

class Solution:
    def primeSubarray(self, nums: List[int], k: int) -> int:
        primes = []
        for i, num in enumerate(nums):
            if num in PRIMES:
                primes.append((i, num))
        n = len(nums)
        i, j, m = 0, 0, len(primes)
        _primes = []
        ans = 0
        while j < m:
            bisect.insort(_primes, primes[j][1])
            while i < j and _primes[-1] - _primes[0] > k:
                _primes.pop(bisect.bisect_left(_primes, primes[i][1]))
                i += 1
            if len(_primes) > 1:
                l = primes[j-1][0]+1 if j-1 == 0 or i == 0 else primes[j-1][0] - primes[i-1][0]
                r = n-primes[j][0] if j == m-1 else primes[j+1][0] - primes[j][0]
                ans += l*r
            j += 1
        return ans

提交代码评测得到：耗时731ms，占用内存27.38MB。