Sum of factor numbers of interval product -- prefix sum idea + fixed one shift two

author ：hans774882968 as well as hans774882968

subject

（sorry No link to the original question found ~） I have an array a, length ≤1e5,1 <= a[i] <= 3. Set interval [l,r] The weight of is the number of factors of the product of interval elements , Find the weight sum of all intervals , model 1e9+7.

Ideas

Interval weight is ( Section 2 The number of + 1) * ( Section 3 The number of + 1), These two quantities can be expressed by prefixes and , Might as well set up s2, s3. Then the demand is
$$\sum _{l=0}^{r-1}(s2[r]-s2[l]+1)\ast (s3[r]-s3[l]+1)$$
Let's enumerate r, Think r Is constant , and l It is changing. . It's easy to take apart
$$r\ast s2[r]\ast s3[r]-s2[r]\ast \sum s3[l]+r\ast s2[r]-s3[r]\ast \sum s2[l]+\sum s2[l]\ast s3[l]-\sum s2[l]+r\ast s3[r]-\sum s3[l]+r$$
except s2, s3 outside , It needs maintenance sum(s2[i]), sum(s3[i]), sum(s2[i]*s3[i]) this 3 Prefixes and arrays .

The code generates a set of random data each time , Multiple runs pass the beat , We just think n^2 Both violence and positive solutions have been correctly realized .

import random

mod = int(1e9) + 7


def data_gen(n):
    return [random.randint(1, 3) for _ in range(n)]


def bf(a):
    n = len(a)
    ans = 0
    for i in range(n):
        v2, v3 = 0, 0
        for j in range(i, n):
            v2 += (a[j] == 2)
            v3 += (a[j] == 3)
            ans = (ans + (v2 + 1) * (v3 + 1) % mod) % mod
    return ans


def solve(a):
    n = len(a)
    s2, s3 = [0], [0]
    for v in a:
        s2.append(s2[-1] + (v == 2))
        s3.append(s3[-1] + (v == 3))
    ss2, ss3, s23 = [0], [0], [0]
    for i in range(1, n + 1):
        ss2.append((ss2[-1] + s2[i]) % mod)
    for i in range(1, n + 1):
        ss3.append((ss3[-1] + s3[i]) % mod)
    for i in range(1, n + 1):
        s23.append((s23[-1] + s2[i] * s3[i] % mod) % mod)
    ans = 0
    for r in range(1, n + 1):
        u = ((r * s2[r] * s3[r] % mod - s2[r] * ss3[r - 1] % mod +
              r * s2[r] % mod - s3[r] * ss2[r - 1] % mod +
              s23[r - 1] - ss2[r - 1] + r * s3[r] % mod -
              ss3[r - 1] + r) % mod + mod) % mod
        ans = (ans + u) % mod
    return ans


if __name__ == '__main__':
    a = [2, 1, 3]
    ans1 = solve(a)
    ans2 = bf(a)
    assert ans1 == 13 and ans2 == 13
    a = data_gen(2000)
    ans1 = solve(a)
    ans2 = bf(a)
    assert ans1 == ans2
    print(a[:20], ans1, ans2)