含最多 K 个可整除元素的子数组

标签: 字典树 数组 哈希表 枚举 哈希函数 滚动哈希

难度: Medium

给你一个整数数组 nums 和两个整数 k 和 p ，找出并返回满足要求的不同的子数组数，要求子数组中最多 k 个可被 p 整除的元素。

如果满足下述条件之一，则认为数组 nums1 和 nums2 是 不同 数组：

  • 两数组长度 不同 ，或者
  • 存在 至少 一个下标 i 满足 nums1[i] != nums2[i] 。

子数组 定义为：数组中的连续元素组成的一个 非空 序列。

示例 1：

输入：nums = [2,3,3,2,2], k = 2, p = 2
输出：11
解释：
位于下标 0、3 和 4 的元素都可以被 p = 2 整除。
共计 11 个不同子数组都满足最多含 k = 2 个可以被 2 整除的元素：
[2]、[2,3]、[2,3,3]、[2,3,3,2]、[3]、[3,3]、[3,3,2]、[3,3,2,2]、[3,2]、[3,2,2] 和 [2,2] 。
注意，尽管子数组 [2] 和 [3] 在 nums 中出现不止一次，但统计时只计数一次。
子数组 [2,3,3,2,2] 不满足条件，因为其中有 3 个元素可以被 2 整除。

示例 2：

输入：nums = [1,2,3,4], k = 4, p = 1
输出：10
解释：
nums 中的所有元素都可以被 p = 1 整除。
此外，nums 中的每个子数组都满足最多 4 个元素可以被 1 整除。
因为所有子数组互不相同，因此满足所有限制条件的子数组总数为 10 。

提示：

  • 1 <= nums.length <= 200
  • 1 <= nums[i], p <= 200
  • 1 <= k <= nums.length

进阶：

你可以设计并实现时间复杂度为 O(n2) 的算法解决此问题吗？

Submission

运行时间: 52 ms

内存: 16.8 MB

class SuffixArray:

    def __init__(self):
        return

    @staticmethod
    def build(s, sig):
        # sa: index is rank and value is pos
        # rk: index if pos and value is rank
        # height: lcp of rank i-th suffix and (i-1)-th suffix
        # sum(height): is count of same substring of s
        # n*(n+1)//2 - sum(height): is count of different substring of s
        # max(height): is the longest duplicate substring
        # sig: number of unique rankings which initially is the size of the character set

        n = len(s)
        sa = list(range(n))
        rk = s[:]
        ll = 0  # ll is the length that has already been sorted, and now it needs to be sorted by 2ll length
        tmp = [0] * n
        while True:
            p = [i for i in range(n - ll, n)] + [x - ll for i, x in enumerate(sa) if x >= ll]
            # for suffixes with a length less than l, their second keyword ranking is definitely
            # the smallest because they are all empty
            # for suffixes of other lengths, suffixes starting at 'sa [i]' rank i-th, and their
            # first ll characters happen to be the second keyword of suffixes starting at 'sa[i] - ll'
            # start cardinality sorting, and first perform statistics on the first keyword
            # first, count how many values each has
            cnt = [0] * sig
            for i in range(n):
                cnt[rk[i]] += 1
            # make a prefix and for easy cardinality sorting
            for i in range(1, sig):
                cnt[i] += cnt[i - 1]

            # then use cardinality sorting to calculate the new sa
            for i in range(n - 1, -1, -1):
                w = rk[p[i]]
                cnt[w] -= 1
                sa[cnt[w]] = p[i]

            # new_sa to check new_rk
            def equal(ii, jj, lll):
                if rk[ii] != rk[jj]:
                    return False
                if ii + lll >= n and jj + lll >= n:
                    return True
                if ii + lll < n and jj + lll < n:
                    return rk[ii + lll] == rk[jj + lll]
                return False

            sig = -1
            for i in range(n):
                tmp[i] = 0

            for i in range(n):
                # compute the lcp
                if i == 0 or not equal(sa[i], sa[i - 1], ll):
                    sig += 1
                tmp[sa[i]] = sig

            for i in range(n):
                rk[i] = tmp[i]
            sig += 1
            if sig == n:
                break
            ll = ll << 1 if ll > 0 else 1

        # height
        k = 0
        height = [0] * n
        for i in range(n):
            if rk[i] > 0:
                j = sa[rk[i] - 1]
                while i + k < n and j + k < n and s[i + k] == s[j + k]:
                    k += 1
                height[rk[i]] = k
                k = 0 if k - 1 < 0 else k - 1
        return sa, rk, height


class Solution:
    def countDistinct(self, nums: List[int], k: int, p: int) -> int:
        cnt = res = 0
        n = len(nums)
        j = 0
        right = [0] * n
        for i in range(n):
            while j < n and cnt + (nums[j] % p == 0) <= k:
                cnt += (nums[j] % p == 0)
                j += 1
            res += j - i
            right[i] = j
            cnt -= (nums[i] % p == 0)
        sa, rk, height = SuffixArray().build(nums, max(nums)+1)
        dup = sum(min(height[i], right[sa[i]] - sa[i]) for i in range(1, n))
        res -= dup
        return res

Explain

这个题解利用了两步主要的算法技巧来解决问题：首先，它使用了双指针技术来计算所有可能的符合条件的子数组的数量。具体来说，外层循环遍历数组，内层循环扩展直到子数组中可以被p整除的元素数量超过k为止。在每次外层循环中，计算从当前索引起始的、满足条件的子数组数量，并逐步减少被p整除的元素的计数。其次，利用后缀数组来识别重复的子数组。构建后缀数组并计算最长公共前缀（LCP），从而可以快速确定哪些子数组是重复的。最后，将找到的总子数组数减去重复的子数组数，得到答案。

时间复杂度: O(n log n)

空间复杂度: O(n)

class SuffixArray:
    def __init__(self):
        return

    @staticmethod
    def build(s, sig):
        n = len(s)
        sa = list(range(n))  # 后缀数组初始化
        rk = s[:]  # 初始化排名数组为字符数组
        ll = 0  # 已排序的长度
        tmp = [0] * n  # 临时数组，用于交换
        while True:
            # 合并排序的核心部分
            p = [i for i in range(n - ll, n)] + [x - ll for i, x in enumerate(sa) if x >= ll]
            cnt = [0] * sig
            for i in range(n):
                cnt[rk[i]] += 1
            for i in range(1, sig):
                cnt[i] += cnt[i - 1]
            for i in range(n - 1, -1, -1):
                w = rk[p[i]]
                cnt[w] -= 1
                sa[cnt[w]] = p[i]
            sig = -1
            for i in range(n):
                if i == 0 or not equal(sa[i], sa[i - 1], ll):
                    sig += 1
                tmp[sa[i]] = sig
            for i in range(n):
                rk[i] = tmp[i]
            sig += 1
            if sig == n:
                break
            ll = ll << 1 if ll > 0 else 1
        k = 0
        height = [0] * n
        for i in range(n):
            if rk[i] > 0:
                j = sa[rk[i] - 1]
                while i + k < n and j + k < n and s[i + k] == s[j + k]:
                    k += 1
                height[rk[i]] = k
                k = 0 if k - 1 < 0 else k - 1
        return sa, rk, height

class Solution:
    def countDistinct(self, nums: List[int], k: int, p: int) -> int:
        cnt = res = 0
        n = len(nums)
        j = 0
        right = [0] * n
        for i in range(n):
            while j < n and cnt + (nums[j] % p == 0) <= k:
                cnt += (nums[j] % p == 0)
                j += 1
            res += j - i
            right[i] = j
            cnt -= (nums[i] % p == 0)
        sa, rk, height = SuffixArray().build(nums, max(nums)+1)
        dup = sum(min(height[i], right[sa[i]] - sa[i]) for i in range(1, n))
        res -= dup
        return res

Explore

在双指针技术中，指针 j 负责扩展子数组的长度直到子数组中可以被 p 整除的元素数量超过 k 为止。为了准确地管理可被 p 整除的元素计数（cnt），每次 j 向右移动一个位置时，首先检查即将包含进子数组的元素 nums[j] 是否能被 p 整除。如果可以，cnt 增加 1。当指针 i 向右移动，即准备考虑下一个新的子数组起点时，需要从 cnt 中减去 nums[i] 的贡献（如果 nums[i] % p == 0，则减 1）。这样，通过逐步调整 cnt，保证了每次内层循环中 cnt 的值都正确地反映了从 i 到 j 子数组中可以被 p 整除的元素数量。

在使用后缀数组和最长公共前缀（LCP）来识别重复的子数组时，首先通过构建后缀数组（suffix array）对所有子数组的起点进行排序。然后，计算相邻后缀（即排序后的子数组）之间的最长公共前缀（LCP）。LCP 数组中的每个值表示排序后相邻的两个子数组的最长公共前缀的长度。为了计算重复的子数组数量，我们考虑后缀数组中的每个元素 sa[i] 和它的右边界 right[sa[i]]（即从 sa[i] 开始的子数组能延伸到的最远位置）。重复子数组的长度至少为 min(height[i], right[sa[i]] - sa[i])，其中 height[i] 是 sa[i] 与 sa[i-1] 的 LCP。通过累加这些值，可以得到所有重复子数组的总长度，从而计算出重复的子数组数量。

在处理整数数组时，不需要将整数转换为字符。排名数组（rank array）初始化时通常是基于数组元素的值来排序。对于整数数组来说，可以直接使用整数的值作为排名。初始排名可以通过将数组元素及其索引位置放入一个列表，然后根据元素的值对此列表进行排序来得到。排序后，元素的新位置（索引）在列表中即代表了它的排名。这种方法避免了将整数转换为字符的步骤，直接利用整数本身的值进行操作，更为直接和高效。