[Leetcode]1513. Number of Substrings With Only 1s

Given a binary string s (a string consisting only of '0' and '1's).
Return the number of substrings with all characters 1's.
Since the answer may be too large, return it modulo 10^9 + 7.
 
Example 1:
Input: s = "0110111" Output: 9 Explanation: There are 9 substring in total with only 1's characters. "1" -> 5 times. "11" -> 3 times. "111" -> 1 time.

It first appeared in Leetcode Weekly Contest 197. The key is to notice that, given a string of '1's with length n, the number of substrings containing all '1's – a_n – is a arithmetic sequence.

With that being said, there are two different ways of implementing it.

Arithmetic sequence sum

We can count all of the substring of '1's and use the arithmetic sequence sum formula.

class Solution:
    def numSub(self, s: str) -> int: 
        ret = l = 0
        for i in range(len(s)):
            if s[i] == '1':
                l += 1
            else:
                ret += l*(l+1)//2
                l = 0
        if l:
            ret += l*(l+1)//2
        return ret % (10**9 + 7)

Accumulatively

Or, credit to @lee215, we can maintain a count for the current number of continuous '1's. When a new '1' appears, there will be count more substrings as shown below.

class Solution:
    def numSub(self, s: str) -> int: 
        ret = l = 0
        for i in range(len(s)):
            if s[i] == '1':
                l += 1
                ret += l
            else:
                l = 0
        return ret % (10**9 + 7)