题目描述

难度分:$1406$

输入$n(1 \leq n \leq 2 \times 10^5)$、$m(1 \leq m \leq 2 \times 10^5)$和长为$n$的数组$a(0 \leq a[i] \leq 10^9)$。

定义$f(b)=sum(b)\% m$。输出$a$的所有非空连续子数组$b$的$f(b)$之和，即$\Sigma_{r=1}^{n}\Sigma_{l=1}^{r}(s[r]-s[l-1]) \% m$，其中$s[i]=\Sigma_{j=1}^{i}a[j]$。

输入样例$1$

3 4
2 5 0

输出样例$1$

10

输入样例$2$

10 100
320 578 244 604 145 839 156 857 556 400

输出样例$2$

2736

算法

树状数组

根据目标函数的表达式$\Sigma_{r=1}^{n}\Sigma_{l=1}^{r}(s[r]-s[l-1]) \% m$，对于一个固定的子数组右端点，我们维护$r=s[i] \% m$，由于取模满足加法分配率，所以此时我们需要知道前面$j \in [1,i]$时$s[j-1] \% m$可以取哪些值？以便分类讨论。

构建两个长度为$m$的树状数组，$tr_1$和$tr_2$，索引$i$就是“前缀和$\% m$”的余数。对于$tr_1$，索引$i$上的值是余数为$i$的“前缀和$\% m$”的累加和；对于$tr_2$，索引$i$上的值是余数为$i$的前缀和个数。初始情况下，$tr_1$所有索引的位置都是$0$，$tr_2$除了$tr_2[0]=1$（空数组前缀和为$0$），所有其他索引上的值都是$0$。这样一来，如果当前$r=s[i] \% m$，那么就要知道前面有多少个$s[l-1] \% m \lt r$，以及这些$s[l-1] \% m$的累加和，假设有$cnt_1$个，累加和为$s_1$；同样的，还需要知道前面有多少个$s[l-1] \% m \gt r$，以及这些$s[l-1] \% m$的累加和，假设有$cnt_2$个，累加和为$s_2$。$cnt_1 \times r-s_1$可以直接减，但是$cnt_2 \times r-s_2 \lt 0$不能直接减，需要对$r$这个余数$+m$才行，相当于计算$(a-b) \% m$时，在$a \geq b$时可以直接化为$a \% m-b \% m$，在$a \lt b$时应该化为$a \% m-b \% m+m$。因此$i$位置对答案的贡献就是$cnt_1 \times r-s_1+cnt_2 \times (r+m)-s_2$。

其中$s_1$、$s_2$、$cnt_1$、$cnt_2$四个值都可以通过对$tr_1$和$tr_2$两个树状数组查询得到。计算完$i$位置的贡献，就把$r$加在$tr_1$的$r$位置，$1$加在$tr_2$的$r$位置。

复杂度分析

时间复杂度

遍历一遍数组$a$就可以求得答案，对于每个$a[i]$，需要对树状数组进行查询和修改操作，时间复杂度为$O(log_2m)$。因此，整个算法的时间复杂度为$O(nlog_2m)$。

空间复杂度

空间消耗的瓶颈就在于树状数组，这两个树状数组都是$O(m)$规模的，因此整个算法的额外空间复杂度为$O(m)$。

C++ 代码

#include <bits/stdc++.h>

using namespace std;
typedef long long LL;
int n, m;

class Fenwick {
public:
    explicit Fenwick(int n): sums_(n + 1) {}

    int lowbit(int x) {
        return x&-x;
    }

    void add(int idx, int val) {
        for(++idx; idx < sums_.size(); idx += lowbit(idx)) {
            sums_[idx] += val;
        }
    }

    LL query(int idx) {
        LL ans = 0;
        for(++idx; idx > 0; idx -= lowbit(idx)) {
            ans += sums_[idx];
        }
        return ans;
    }

    LL query(int left, int right) {
        if(left > right) return 0LL;
        return query(right) - query(left - 1);
    }

private:
    vector<LL> sums_;
};

int main() {
    scanf("%d%d", &n, &m);
    Fenwick tr1(m), tr2(m);
    LL ans = 0, s = 0;
    tr2.add(0, 1);
    for(int i = 1; i <= n; i++) {
        int a;
        scanf("%d", &a);
        s += a;
        LL r = s % m;
        LL cnt1 = tr2.query(0, r - 1), cnt2 = tr2.query(r + 1, m - 1);
        ans += cnt1 * r - tr1.query(0, r - 1) + cnt2 * (r + m) - tr1.query(r + 1, m - 1);
        tr1.add(r, r);
        tr2.add(r, 1);
    }
    printf("%lld\n", ans);
    return 0;
}