LC 3261. 统计满足 K 约束的子字符串数量 II
https://leetcode.cn/problems/count-substrings-that-satisfy-k-constraint-ii/description/
这题想了挺久的，有好几个点都是慢慢想到的：
- 对于每个位置i来说，都有一个 `end[i]` 对应，表示从i开始满足约束的话，最长到达的位置。
- 这个end是个单调非递减的数组，所以可以使用二分搜索。
- 对于query (a, b) 来说，可以分为两个区间处理：
- `[a, p]` 在这个区间里面的位置，他们的end不超过b. 所以满足这些的约束子串可以使用前缀和计算出来。
- `[p+1], b]` 在这个区间里面的位置，他们的结尾是 `b`.
class Solution:
def countKConstraintSubstrings(self, s: str, k: int, queries: List[List[int]]) -> List[int]:
n = len(s)
def build_end():
a, b = n, n
right = [n] * (n + 1)
for i in reversed(range(n)):
if s[i] == '0':
right[i] = a
a = i
else:
right[i] = b
b = i
a, b = n, n
for i in range(n):
if s[i] == '0':
a = i
break
for i in range(n):
if s[i] == '1':
b = i
break
for _ in range(k):
a = right[a]
b = right[b]
end = [-1] * n
for i in range(n):
end[i] = max(a, b)
if s[i] == '0':
a = right[a]
else:
b = right[b]
return end
end = build_end()
def search(x):
s, e = 0, n - 1
while s <= e:
m = (s + e) // 2
if end[m] <= x:
s = m + 1
else:
e = m - 1
return e
acc = [0] * (n + 1)
for i in range(n):
acc[i + 1] = acc[i] + (end[i] - i)
# print(end)
ans = []
for a, b in queries:
p = search(b + 1)
r = 0
# [a...p]
# where end[p] <= (b+1)
if p >= a:
r += acc[p + 1] - acc[a]
else:
p = a - 1
# [p + 1 .. end + 1]
# [p + 2 .. end + 1]
# [end + 1.. end + 1]
sz = (b + 1 - p)
r += (b - p) * sz - (sz - 1) * sz // 2
ans.append(r)
return ans