所有子串的权重总和不超过K

简介

给定一个长度为n的字符串s和权重数组w，定义s的一个子串为s[left:right+1]（其中0<=left<=right<n）。给每个子串分配权重，即s[left:right+1]的权重为w[left]+w[left+1]+…+w[right]。现在需要找到所有权重之和不超过K的子串个数。

解法

暴力枚举

我们可以使用暴力枚举的方法来找到所有子串，并计算它们的权重之和。时间复杂度为$O(n^3)$，会超时。

def count_substrings(s: str, w: List[int], k: int) -> int:
    n = len(s)
    res = 0
    for left in range(n):
        for right in range(left, n):
            sum_w = sum(w[left:right+1])
            if sum_w <= k:
                res += 1
    return res

双指针

我们可以使用双指针来优化暴力枚举。时间复杂度为$O(n^2)$，依然会超时。

def count_substrings(s: str, w: List[int], k: int) -> int:
    n = len(s)
    res = 0
    for left in range(n):
        sum_w = 0
        right = left
        while right < n:
            sum_w += w[right]
            if sum_w <= k:
                res += 1
            else:
                break
            right += 1
    return res

前缀和+二分

我们可以先对权重数组求出前缀和，然后对于每一个右端点，找到满足条件的最小左端点。时间复杂度为$O(nlogn)$。

def count_substrings(s: str, w: List[int], k: int) -> int:
    n = len(s)
    pre_sum = [0] * (n + 1)
    for i in range(1, n + 1):
        pre_sum[i] = pre_sum[i - 1] + w[i - 1]
    res = 0
    for right in range(n):
        left = bisect_left(pre_sum, pre_sum[right + 1] - k)
        res += right - left + 1
    return res

其中，bisect_left是python标准库中的二分查找函数。

滑动窗口

我们可以使用滑动窗口来优化前缀和+二分的方法。滑动窗口维护一个区间，使得区间内所有子串的权重之和不超过K。时间复杂度为$O(n)$。

def count_substrings(s: str, w: List[int], k: int) -> int:
    n = len(s)
    left = right = sum_w = res = 0
    while right < n:
        sum_w += w[right]
        while sum_w > k:
            sum_w -= w[left]
            left += 1
        res += right - left + 1
        right += 1
    return res

总结

本文介绍了四种方法来求所有子串的权重总和不超过K的个数。暴力枚举和双指针时间复杂度为$O(n^3)$和$O(n^2)$，不适合大规模数据。前缀和+二分和滑动窗口时间复杂度为$O(nlogn)$和$O(n)$，其中滑动窗口方法最优。