Fruit Sequences (Codeforces 1428F)
Tags: DP, divide and conquer

Zookeeper is buying a carton of fruit to feed his pet wabbit. The fruits are a sequence of apples and oranges, which is represented by a binary string $$$s_1s_2\ldots s_n$$$ of length $$$n$$$. $$$1$$$ represents an apple and $$$0$$$ represents an orange.Since wabbit is allergic to eating oranges, Zookeeper would like to find the longest contiguous sequence of apples. Let $$$f(l,r)$$$ be the longest contiguous sequence of apples in the substring $$$s_{l}s_{l+1}\ldots s_{r}$$$. Help Zookeeper find $$$\sum_{l=1}^{n} \sum_{r=l}^{n} f(l,r)$$$, or the sum of $$$f$$$ across all substrings.

Input: The first line contains a single integer $$$n$$$ $$$(1 \leq n \leq 5 \cdot 10^5)$$$. The next line contains a binary string $$$s$$$ of length $$$n$$$ $$$(s_i \in \{0,1\})$$$ 

Output: Print a single integer: $$$\sum_{l=1}^{n} \sum_{r=l}^{n} f(l,r)$$$. 

Note: In the first test, there are ten substrings. The list of them (we let $$$[l,r]$$$ be the substring $$$s_l s_{l+1} \ldots s_r$$$):  $$$[1,1]$$$: 0  $$$[1,2]$$$: 01  $$$[1,3]$$$: 011  $$$[1,4]$$$: 0110  $$$[2,2]$$$: 1  $$$[2,3]$$$: 11  $$$[2,4]$$$: 110  $$$[3,3]$$$: 1  $$$[3,4]$$$: 10  $$$[4,4]$$$: 0 The lengths of the longest contiguous sequence of ones in each of these ten substrings are $$$0,1,2,2,1,2,2,1,1,0$$$ respectively. Hence, the answer is $$$0+1+2+2+1+2+2+1+1+0 = 12$$$.

Sample Inputs:
4
0110

7
1101001

12
011100011100

Sample Outputs:
12

30

156

Solution (Python):
hist = [0] * 1000005

def solve(n, s):
    tot = 0
    cur = 0
    i = 0
    while i < n:
        if s[i] == '0':
            tot += cur
        else:
            l = i
            r = i
            while r + 1 < n and s[r + 1] == '1':
                r += 1
            for x in range(r - l + 1):
                cur += l + x + 1 - hist[x]
                tot += cur
                hist[x] = r - x + 1
            i = r
        i += 1
    return tot
n = int(input())
s = input()
print(solve(n, s))

Solution (Python):
N = int(input())
S = input()
p = 0
b = [0 for i in range(N + 1)]
a = 0
s = 0
for c in S:
    if c == '0':
        p = 0
        b[0] += 1
    elif p == 0:
        p = 1
        s += b[0] + 1
        b[1] += b[0]
        b[0] = 0
        b[1] += 1
    else:
        p += 1
        s += b[p - 1] + p - 1
        b[p] += b[p - 1]
        b[p - 1] = 1
    a += s
print(a)

Solution (Python):
hist = [0] * 1000005

def solve(n, s):
    tot = 0
    cur = 0
    i = 0
    while i < n:
        if s[i] == '0':
            tot += cur
        else:
            l = i
            r = i
            while r + 1 < n and s[r + 1] == '1':
                r += 1
            for x in range(r - l + 1):
                cur += l + x + 1 - hist[x]
                tot += cur
                hist[x] = r - x + 1
            i = r
        i += 1
    return tot
n = int(input())
s = input()
print(solve(n, s))