Number_Theory/Mf_Sieve.hpp

View this file on GitHub
Last update: 2023-02-11 22:28:05+08:00
Include: #include "Number_Theory/Mf_Sieve.hpp"

Depends on

Verified with

Verify/Sum_of_Totient_Function.test.cpp

Code

#pragma once

#include "Prime_Sieve.hpp"

template <class T>
std::vector<T> mf_sieve(int n, std::function<T(i64, i64, i64)> f) {
    /*
        ##pragma f is a multiplicative-function
        f(n,p,c) <-> n=p^c
    */
    std::vector<T> ans(n + 1, T());
    std::vector<int> ps = prime_sieve(n);
    int p(ps.size());

    std::function<void(int, i64, T)> dfs = [&](int i, i64 x, T y) {
        ans[x] = y;
        if (y == T()) return;
        for (int j = i + 1; j < p; j++) {
            i64 nx = x * ps[j];
            i64 dx = ps[j];
            if (nx > n) break;
            for (int c = 1; nx <= n; nx *= ps[j], dx *= ps[j], ++c)
                dfs(j, nx, y * f(dx, ps[j], c));
        }
    };
    ans[1] = 1;
    dfs(-1, 1, 1);
    return ans;
};

#line 2 "Number_Theory/Mf_Sieve.hpp"

#line 2 "Number_Theory/Prime_Sieve.hpp"

#line 2 "Template/Template.hpp"

#include <bits/stdc++.h>

using i64 = std::int64_t;
#line 4 "Number_Theory/Prime_Sieve.hpp"

std::vector<int> prime_sieve(int N) {
    std::vector<bool> sieve(N / 3 + 1, 1);
    for (int p = 5, d = 4, i = 1, sqn = sqrt(N); p <= sqn;
         p += d = 6 - d, i++) {
        if (!sieve[i]) continue;
        for (int q = p * p / 3, r = d * p / 3 + (d * p % 3 == 2), s = 2 * p,
                 qe = sieve.size();
             q < qe; q += r = s - r)
            sieve[q] = 0;
    }
    std::vector<int> ret{2, 3};
    for (int p = 5, d = 4, i = 1; p <= N; p += d = 6 - d, i++)
        if (sieve[i]) ret.push_back(p);
    while (!ret.empty() && ret.back() > N) ret.pop_back();
    return ret;
}
#line 4 "Number_Theory/Mf_Sieve.hpp"

template <class T>
std::vector<T> mf_sieve(int n, std::function<T(i64, i64, i64)> f) {
    /*
        ##pragma f is a multiplicative-function
        f(n,p,c) <-> n=p^c
    */
    std::vector<T> ans(n + 1, T());
    std::vector<int> ps = prime_sieve(n);
    int p(ps.size());

    std::function<void(int, i64, T)> dfs = [&](int i, i64 x, T y) {
        ans[x] = y;
        if (y == T()) return;
        for (int j = i + 1; j < p; j++) {
            i64 nx = x * ps[j];
            i64 dx = ps[j];
            if (nx > n) break;
            for (int c = 1; nx <= n; nx *= ps[j], dx *= ps[j], ++c)
                dfs(j, nx, y * f(dx, ps[j], c));
        }
    };
    ans[1] = 1;
    dfs(-1, 1, 1);
    return ans;
};