4

题目

游游有一个只包含'0'和'1'的字符串，他想知道这个字符串有多少个好子串？

一个字符串如果是“好串”，那么该字符串的所有前缀，'0'的数量严格大于'1'的数量。

• 输入描述

输入一个只包含'0'和'1'的字符串，长度不超过100000。

• 输出描述

输出一个整数，代表答案。

• 示例1

输入

100

输出

3

分析

熟练度不太行，想到一点思路没来得及写出来...

首先看数据范围，n <= 1e5，所以暴力遍历每个区间必然会超时。我们使用的算法的时间复杂度需要是 O(nlogn) 或者 O(n)。

由定义可以得到，如果一个子串是“好串”，那么依次从它的尾部删去一个字符，剩余的子串依然是“好串”。所以对于每个位置，如果以该位置为开头的最长“好串”长度为 length，那么以该位置为开头的“好串”就有 length 个。

如果把字符串的 '0' 对应到 1，'1' 对应到 -1，那么 [i, j] 为好串的条件就等价于 i 处为 '0' 且前缀和数组 pre 对所有 i<=k<=j 满足 pre[k] >= pre[i]。寻找以 i （下标从1开始）为开头的最长“好串”，就等价于寻找 i 右侧第一个使得 [i,j] 不是“好串”的最小的 j。这等价于寻找 i 右侧使得 pre[j] < pre[i] 的最小的 j，问题转化为典型的单调栈问题。

代码

s = input()
n = len(s)
pre = [0] * (n + 1)
str2num = {"0": 1, "1": -1}
for i in range(1, n + 1):
    pre[i] = pre[i - 1] + str2num[s[i - 1]]

stack = [(n + 1, int(1e9))] # (idx, pre[idx]) 初始化的时候加入最右侧的哨兵
cnt = 0
i = n
j = n # 单调栈未处理过的前缀和数组的最右侧的位置
while i > 0:
    # 对每个 i 找到 i 右侧第一个小于 pre[i] 的 pre[j] 和 j
    if s[i - 1] == "0":
        while j > i:
            while stack and stack[-1][1] >= pre[j]:
                stack.pop()
            stack.append((j, pre[j]))
            j -= 1
        if stack[0][1] >= pre[i]: # 右侧最小值
            cnt += n + 1 - i
        else:
            while stack[-1][1] >= pre[i]:
                stack.pop()
            cnt += stack[-1][0] - i
    i -= 1
print(cnt)

i 和 j 只会向左移动，前缀和数组的每项最多入栈一次、出栈一次，时间复杂度为 O(n)。