牛客练习赛81 E 小Q与函数求和1

其中是完全积性函数,可线性筛。
其余部分可以处理。

#include <cstdio>
#include <algorithm>
#include <cstring>
using namespace std;

typedef long long ll;

const int N = 5e6 + 100;
const int MOD = 998244353;

void upd(int &a, int b) {
    a += b;
    if (a >= MOD) a -= MOD;
}

int pri[N], phi[N], mu[N], fk[N], rev[N], s[N], f[N], tot, n, k;
bool is_pri[N];

ll qpow(ll x, ll n) {
    ll res = 1;
    while (n > 0) {
        if (n & 1) res = res * x % MOD;
        n /= 2;
        x = x * x % MOD;
    }
    return res;
}

void init() {
    memset(is_pri, true, sizeof(is_pri));
    mu[1] = phi[1] = fk[1] = 1;
    for (int i = 2; i <= n; i++) {
        if (is_pri[i]) {
            pri[tot++] = i;
            phi[i] = i - 1;
            fk[i] = qpow(i, k + 1);
            mu[i] = MOD - 1;
        }
        for (int j = 0; i * pri[j] <= n; j++) {
            is_pri[i * pri[j]] = false;
            fk[i * pri[j]] = 1LL * fk[i] * fk[pri[j]] % MOD;
            if (i % pri[j] == 0) {
                phi[i * pri[j]] = phi[i] * pri[j];
                mu[i * pri[j]] = 0;
                break;
            }
            else {
                phi[i * pri[j]] = phi[i] * phi[pri[j]];
                mu[i * pri[j]] = MOD - mu[i];
            }
        }
    }

    rev[0] = rev[1] = 1;
    for (int i = 2; i <= n; i++) rev[i] = 1LL * (MOD - MOD / i) * rev[MOD % i] % MOD;
    for (int i = 1; i <= n; i++) {
        for (int j = i; j <= n; j += i) {
            upd(s[i], phi[j]);
            if (mu[j / i]) upd(f[j], 1LL * fk[i] * rev[phi[i]] % MOD * mu[j / i] % MOD);
        }
    }
}

int main() {
    //freopen("0.txt", "r", stdin);
    scanf("%d%d", &n, &k); init();
    int ans = 0;
    for (int i = 1; i <= n; i++) ans = (ans + 1LL * f[i] * s[i] % MOD * s[i]) % MOD;
    ans = (ans % MOD + MOD) % MOD;
    printf("%d\n", ans);
    return 0;
}