ACM-ICPC 2018 沈阳赛区网络预赛 G. Spare Tire (容斥原理)

可推出$a_n = n^2+n, $ 设\(S_n = \sum_{i=1}^{n} a_i\) 则 \(S_n = \frac{n(n+1)(2n+1)}{6} + \frac{n(n+1)}{2}\)
需要求出\([1,N]\)中与\(M\)互质的下标的和
可以容斥计算答案,\(O(1)\)时间算出\(S_n\),需要减去与\(M\)非互质的下标\(a_i\)
对于\(M\)的每一种质因数的组合\(k\),可以求出\(b_kn = (kn)^2+kn\),则也可以\(O(1)\)得到在\(bn在[1-N]\)中的和.
根据容斥原理,奇数个质因子的组合要减去,偶数个的加回

#include <iostream>
#include <string>
#include <cmath>
#include <stdio.h>
#include <vector>
using namespace std;
const int mod=  1e9+7;
typedef long long LL;
vector<LL> fac;

LL qpow(LL a,LL n)
{
    LL res=1;
    while(n){
        if(n&1) res = res*a %mod;
        a = a*a %mod;
        n>>=1;
    }
    return res;
}

LL solve(LL n,LL m)
{
    LL rev6 = qpow((LL)6,mod-2);
    LL rev2 = qpow((LL)2,mod-2);
    fac.clear();
    LL tmp = m;
    for(LL i=2;i*i<=tmp;++i){
        if(tmp%i==0){
            fac.push_back(i);
            while(tmp%i==0) tmp/=i;
        }
    }
    if(tmp>1) fac.push_back(tmp);

    LL res = n *(n+1) %mod *(2*n+1) %mod *rev6 %mod;
    LL t = n*(n+1) %mod *rev2 %mod;
    res = (res+t) %mod;

    int len = fac.size();
    int tot = 1<<len;
    for(int i=1;i<tot;++i){
        int cnt = __builtin_popcount(i);
        LL tmp =1;
        for(int j = 0;j<len;++j){
            if(i&(1<<j)){
                tmp *= fac[j];
            }
        }
        LL nn = n/tmp;
        LL pt = (nn)%mod *(nn+1)%mod *(2*nn+1) %mod*rev6%mod;
        pt = pt*tmp %mod *tmp %mod;
        pt = (pt+nn*(tmp+tmp*nn)%mod*rev2%mod) %mod;
        if(cnt&1) res  = (res-pt+mod)%mod;
        else res = (res+pt)%mod;
    }
    return res;
}

int main()
{
    #ifndef ONLINE_JUDGE
        freopen("in.txt","r",stdin);
        freopen("out.txt","w",stdout);
    #endif
    LL N,M;
    while(scanf("%lld %lld",&N,&M)==2){
        LL res = solve(N,M);
        printf("%lld\n",res);
    }
    return 0;
}