P10740 [SEERC 2020] Divisible by 3 题解

宗傂轥

已于 2025-06-04 13:52:45 修改

阅读量1.3k

点赞数 26

文章标签： c++ 数据结构算法动态规划

于 2025-04-04 15:42:39 首次发布

本文链接：https://blog.youkuaiyun.com/2401_89520240/article/details/146998411

版权

题目描述

定义一个序列 b 的权重为 $\sum_{i=1}^n \sum_{j=1}^{i-1} b_i \times b_j$ 。

现在你有一个长度为 n 的数组 a，求一共存在多少种 (l,r) 使得 1≤l≤r≤n 且 $[a_l,a_{l+1},$ … $,a_r]$ 的权重能被 3 整除。

输入格式

第一行一个整数 n (1 ≤ n ≤ 5× $10^5$ )。

然后 n 个整数 ai (0 ≤ $a_i$ ≤ $10^9$ )。

输出格式

输出方案总数。

输入输出样例

输入 #1

3
5 23 2021

输出 #1

输入 #2

5
0 0 1 3 3

输出 #2

输入 #3

10
0 1 2 3 4 5 6 7 8 9

输出 #3

说明/提示

对于第一个样例，存在 [1,1]、[2,2]、[3,3]、[1,3] 共 4 种方案。

方法思路

问题转化：首先，权重的计算可以转化为数学公式。对于一个子数组 $b = [a_l, a_{l+1}, \ldots, a_r]$ ，其权重 W 可以表示为：

因此， $W \equiv 0 \pmod{3}$ 等价于 $\left(S^2 - \sum_{i=1}^n a_i^2\right) \equiv 0 \pmod{6}$ ，其中 S 是子数组的和。由于 2 和 3 互质，这可以分解为两个条件：
- $S^2 \equiv \sum a_i^2 \mod 3$
- 由于模3运算中，平方后的可能值为0或1（因为1² ≡ 1 mod 3，2² ≡ 1 mod 3），所以：
  - 如果 S ≡ 0 mod 3，则 $0 \equiv \sum a_i^2 \mod 3 \Rightarrow \sum a_i^2 \equiv 0 \mod 3$ 。
  - 如果 S ≡ 1或2 mod 3，则 $1 \equiv \sum a_i^2 \bmod 3 \Rightarrow \sum a_i^2 \equiv 1 \bmod 3$
关键观察：我们需要同时跟踪子数组的和 S mod 3 和子数组的平方和 mod 3。具体来说，对于每个位置，维护三个状态：
- $\left(S \bmod 3, \sum a_i^2 \bmod 3\right)$ 的可能组合：(0,0), (0,1), (0,2), (1,0), (1,1), (1,2), (2,0), (2,1), (2,2)。
- 其中，满足条件的是 (0,0)、(1,1) 和 (2,1)。
前缀和优化：使用前缀和数组来高效计算子数组的和和平方和。具体来说：
- 前缀和数组 $\text{sum}[i] = (\text{sum}[i-1] + a[i]) \bmod 3$ 。
- 前缀平方和数组 $\text{sqsum}[i] = (\text{sqsum}[i-1] + a[i]^2) \bmod 3$ 。
- 对于任意子数组 [l, r]，其和 mod 3 为 $(sum[r] - sum[l - 1]) \bmod 3$ ，平方和 mod 3 为 $(sqsum[r] - sqsum[l-1]) \bmod 3$ 。
哈希统计：我们需要统计满足特定条件的前缀对的数量。具体来说，维护一个字典或数组 count[s][q]count[s][q] 记录之前出现过的前缀和 mod 3 和前缀平方和 mod 3 的组合 (s, q) 的次数。遍历数组时，对于当前位置 i，计算当前的前缀和 s 和前缀平方和 q，然后查找之前有多少个位置 j 满足：
- $(s - s_j) \, mod \, 3 = 0$ 且 $(q - q_j) \,mod\, 3 = 0$
- $(s - s_j) \bmod 3 = 1$ 且 $(q - q_j) \, mod \,{3} \equiv 1$ 。
- $(s - s_j) \bmod 3 = 2$ 且 $(q - q_j) \bmod 3 = 1$ 。
  这些 j 的数量之和即为以 i 结尾的满足条件的子数组数量。

代码

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int MAXN = 1e6+10;

int a[MAXN], n;
int cnt[3][3];

signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        a[i] %= 3;
    }

    int res = 0, sum = 0, sqsum = 0;
    cnt[0][0] = 1;

    for (int i = 1; i <= n; i++) {
        int num = a[i];
        sum = (sum + num) % 3;
        int num_sq = (num * num) % 3;
        sqsum = (sqsum + num_sq) % 3;

        for (int ds = 0; ds < 3; ds++) {
            int dq;
            if (ds == 0) {
                dq = 0;
            } else {
                dq = 1;
            }
            int s = (sum - ds + 3) % 3;
            int q = (sqsum - dq + 3) % 3;
            res += cnt[s][q];
        }

        cnt[sum][sqsum]++;
    }

    cout << res << '\n';
    return 0;
}

解说

全局变量与常量：
- int a[MAXN]：存储输入数组及其长度。
- int cnt[3][3]：二维数组，记录前缀和的模3状态及其平方和的模3状态的出现次数。
主函数：
- 读取输入数组并取模3。
- 初始化计数器cnt[0][0] = 1，处理空前缀的情况。
- 遍历数组，维护当前前缀和（sum）与平方和（sqsum）的模3值。
- 对每个可能的历史前缀和模3值（ds）计算对应的子数组和（s）与平方和（q），并累加符合条件的计数。
- 更新计数器cnt。

核心逻辑

条件判断：
- 子数组的和s = (sum_i - sum_j) % 3。
- 子数组的平方和q = (sqsum_i - dq) % 3（dq为历史前缀和平方的模3值）。
- 需要满足s² ≡ q (mod 3)。
前缀和优化：
- 遍历每个元素时，利用前缀和技巧，通过历史状态快速统计满足条件的子数组数量。
- cnt数组记录所有可能的前缀和状态组合，避免重复计算。