2025-03-30：统计满足 K 约束的子字符串数量Ⅱ。用go语言，给定一个二进制字符串 s 和一个整数 k，还有一个二维整数

福大大架构师每日一题

发布于 2025-03-31 19:43:48

3500

代码可运行

文章被收录于专栏：福大大架构师每日一题福大大架构师每日一题

运行总次数：0

代码可运行

2025-03-30：统计满足 K 约束的子字符串数量Ⅱ。用go语言，给定一个二进制字符串 s 和一个整数 k，还有一个二维整数数组 queries，其中每个元素 queries[i] = [li, ri] 代表一个查询。

我们定义一个二进制字符串满足 k 约束，条件是：

1.字符串中 0 的数量不能超过 k。

2.字符串中 1 的数量不能超过 k。

要求返回一个整数数组 answer，其中 answer[i] 表示在子字符串 s[li..ri] 中满足 k 约束的子字符串数量。

1 <= s.length <= 100000。

s[i] 是 '0' 或 '1'。

1 <= k <= s.length。

1 <= queries.length <= 100000。

queries[i] == [li, ri]。

0 <= li <= ri < s.length。

所有查询互不相同。

输入：s = "0001111", k = 2, queries = [[0,6]]。

输出：[26]。

解释：

对于查询 [0, 6]， s[0..6] = "0001111" 的所有子字符串中，除 s[0..5] = "000111" 和 s[0..6] = "0001111" 外，其余子字符串都满足 k 约束。

题目来自leetcode3261。

大体步骤如下：

1. 首先定义一个函数 countKConstraintSubstrings，接受一个二进制字符串 s，一个整数 k，和一个二维整数数组 queries，并返回一个整数数组作为结果（满足 k 约束的子字符串数量）。
2. 在函数内部，初始化变量 count 用来统计当前窗口中 0 和 1 的数量，初始化 right 数组用来存放每个位置右侧第一个不满足 k 约束的位置，初始化 prefix 数组用来存放前缀和。
3. 遍历字符串 s，同时维护窗口，根据题目要求更新 count，当某一位超过 k 时移动左指针 i，更新 right 数组并计算该窗口的前缀和存入 prefix 数组。
4. 遍历 queries，对每个查询 [li, ri] 找到满足 k 约束的子字符串的数量。首先找到右边界 i，然后计算子字符串中满足 k 约束的数量，一部分为在 i 左侧的数量，另一部分为前缀和差值，将结果存入 res 数组。
5. 返回 res 数组作为结果。

总的时间复杂度为 O(n + m)，其中 n 是字符串的长度，m 是查询数组 queries 的长度。

额外空间复杂度为 O(n + m)，用来存储 right 数组和 prefix 数组。

Go完整代码如下：

package main

import (
    "fmt"
)

func countKConstraintSubstrings(s string, k int, queries [][]int) []int64 {
    n := len(s)
    count := [2]int{}
    right := make([]int, n)
    for i := range right {
        right[i] = n
    }
    prefix := make([]int64, n+1)
    i := 0
    for j := 0; j < n; j++ {
        count[int(s[j]-'0')]++
        for count[0] > k && count[1] > k {
            count[int(s[i]-'0')]--
            right[i] = j
            i++
        }
        prefix[j+1] = prefix[j] + int64(j-i+1)
    }

    res := make([]int64, 0, len(queries))
    for _, query := range queries {
        l, r := query[0], query[1]
        i := min(right[l], r+1)
        part1 := int64(i-l+1) * int64(i-l) / 2
        part2 := prefix[r+1] - prefix[i]
        res = append(res, part1+part2)
    }
    return res
}

func min(a, b int)int {
    if a < b {
        return a
    }
    return b
}

func main() {
    s := "0001111"
    k := 2
    queries := [][]int{{0,6}}
    result := countKConstraintSubstrings(s, k,queries)
    fmt.Println(result)
}

Python完整代码如下：

# -*-coding:utf-8-*-

defcount_k_constraint_substrings(s, k, queries):
    n = len(s)
    count = [0, 0]
    right = [n] * n
    prefix = [0] * (n + 1)
    i = 0

    for j inrange(n):
        count[int(s[j])] += 1
        while count[0] > k and count[1] > k:
            count[int(s[i])] -= 1
            right[i] = j
            i += 1
        prefix[j + 1] = prefix[j] + (j - i + 1)

    res = []
    for l, r in queries:
        i = min(right[l], r + 1)
        part1 = (i - l + 1) * (i - l) // 2
        part2 = prefix[r + 1] - prefix[i]
        res.append(part1 + part2)

    return res

# 示例
s = "0001111"
k = 2
queries = [[0, 6]]
result = count_k_constraint_substrings(s, k, queries)
print(result)