区间乘积的因子数之和——前缀和思想+定一移二

Posted 2022-06-27 hans774882968

文章目录

• • • 题目
    • 思路

作者：hans774882968以及hans774882968

题目

（sorry没找到原题链接~）有一个数组a，长度≤1e5，1 <= a[i] <= 3。设区间[l,r]的权值为区间元素乘积的因数个数，求所有区间的权值和，模1e9+7。

思路

区间权值为(区间2的个数 + 1) * (区间3的个数 + 1)，这两个量可以用前缀和表示，不妨设为s2, s3。则所求为
$$\sum _{l=0}^{r-1}(s2[r]-s2[l]+1)\ast (s3[r]-s3[l]+1)$$
我们枚举r，则认为r是固定的，而l是变化的。拆开得
$$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$$
除了s2, s3之外，还需要维护sum(s2[i]), sum(s3[i]), sum(s2[i]*s3[i])这3个前缀和数组。

代码每次会生成一组随机数据，多次运行都对拍通过，我们就认为n^2暴力和正解都正确实现了。

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)