Leetcode 3234. Count the Number of Substrings With Dominant Ones

• Leetcode 3234. Count the Number of Substrings With Dominant Ones
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3234. Count the Number of Substrings With Dominant Ones

1. 解题思路

这一题要求被1 dominate的substring的个数，整体上来说是一个$O(N^2)$算法复杂度的东西，还挺麻烦的，所幸题中限制了1的个数需要多于0的平方数，因此我们可以反其道而行，依次考察所有存在$i$个零的substring，看看其中有多少个满足1的个数不少于$i^2$，将其累加即为我们最终的答案，如此一来，整体的算法复杂度可以优化至$O(N^{3/2})$，勉强可以接受吧。

于是，我们只需要给出所有0所在的位置，然后考察所有恰好存在$i$个0的区间，然后看看有多少个满足至少存在$i^2$个1。假设对应的窗口之间有$m$个1，然后区间左右分别有$l,r$个1，此时，我们不难通过容斥原理得到对应的数学解答如下：

$$N=(\sum _{i=1}^{l+r-(i^2-m)+1}i)-(\sum _{i=1}^{l-(i^2-m)+1}i)+\max (0,l-(i^2-m)+1)-(\sum _{i=1}^{r-(i^2-m)+1}i)+\max (0,r-(i^2-m)+1)$$

而等差数列的求解又可以直接写作$\sum _{i=1}^n=\frac{n(n+1)}{2}$，因此，我们不难在$O(1)$的时间复杂度范围内直接获得对应的答案。

综上，总的时间复杂度就是$O(N^{3/2})$，不算很好，但是勉强够用吧……

2. 代码实现

给出python代码实现如下：

class Solution:
    def numberOfSubstrings(self, s: str) -> int:
        n = len(s)
        zeros = [i for i, ch in enumerate(s) if ch == "0"]
        if len(zeros) <= 1:
            return n * (n+1) // 2 - len(zeros)
        
        m = len(zeros)
        ans = (n-1 - zeros[-1]) * (n - zeros[-1]) // 2
        for i in range(m):
            if i == 0:
                one = zeros[i]
            else:
                one = zeros[i] - zeros[i-1]-1
            ans += one*(one+1) // 2
        for i in range(1, m+1):
            if i*i+i > n:
                break
            for j in range(m+1-i):
                l, r = zeros[j], zeros[i+j-1]
                mid = r-l+1-i
                l1 = l if j == 0 else zeros[j] - zeros[j-1]-1
                r1 = n-1-r if i+j-1 == m-1 else zeros[i+j] - zeros[i+j-1]-1
                if mid >= i*i:
                    need = 0
                    cnt = (l1+1) * (r1+1)
                elif mid + l1 + r1 >= i*i:
                    need = i*i-mid
                    cnt = (l1+r1-need+1) * (l1+r1-need+2) // 2
                    if l1 >= need:
                        cnt = cnt - ((l1-need+1) * (l1-need+2) // 2) + (l1-need+1)
                    if r1 >= need:
                        cnt = cnt - ((r1-need+1) * (r1-need+2) // 2) + (r1-need+1)
                else:
                    need = 0
                    cnt = 0
                ans += cnt
        return ans      

提交代码评测得到：耗时7478ms，占用内存18.4MB。