Math/multiplicative_function.h

• View this file on GitHub
• Last update: 2022-12-27 15:34:50+08:00
• Link: https://cses.fi/problemset/task/1713

Verified with

• Math/tests/cnt_divisors_stress.test.cpp

Code

// NOTE: calculate upto N-1
//
// Multiplicative function {{{
template<int N>
struct MultiplicativeFunction {
    // Init sieve and pk
    MultiplicativeFunction() {
        // Init sieve
        for (int i = 2; i*i < N; i++) {
            if (!sieve[i]) {
                for (int j = i*i; j < N; j += i) {
                    sieve[j] = i;
                }
            }
        }

        // Init pk
        for (int i = 2; i < N; i++) {
            if (!sieve[i]) {
                pk[i] = {i, 1};
            } else {
                int p = sieve[i];

                if (pk[i/p].first == p) { // i = p^k
                    pk[i] = {p, pk[i/p].second + 1};
                } else {
                    pk[i] = {-1, 0};
                }
            }
        }
    }

    // Tested: https://cses.fi/problemset/task/1713
    array<int, N> divisors() {
        array<int, N> res;
        res[1] = 1;

        for (int i = 2; i < N; i++) {
            if (pk[i].first > 0) {  // i = p^k
                res[i] = pk[i].second + 1;
            } else {
                // i = u * v, gcd(u, v) = 1
                int u = i, v = 1;
                int p = sieve[i];
                while (u % p == 0) {
                    u /= p;
                    v *= p;
                }
                res[i] = res[u] * res[v];
            }
        }

        return res;
    }

    // mobius(n) = 1  if n is square-free and has *even* number of prime factors
    // mobius(n) = -1 if n is square-free and has *odd* number of of prime factors
    // mobius(n) = 0  if n is not square-free
    array<int, N> mobius() {
        array<int, N> res;
        res[1] = 1;

        for (int i = 2; i < N; ++i) {
            if (pk[i].first > 0) {  // i = p^k
                res[i] = (pk[i].second >= 2) ? 0 : -1;
            } else {
                // i = u * v, gcd(u, v) = 1
                int u = i, v = 1;
                int p = sieve[i];
                while (u % p == 0) {
                    u /= p;
                    v *= p;
                }
                res[i] = res[u] * res[v];
            }
        }
        return res;
    }

// private:
    // sieve[i] == 0 if i is prime,
    // sieve[i] = any prime factor p otherwise
    array<int, N> sieve = {0};

    // pk[i] = {p, k} if i == p^k
    // pk[i] = {-1, 0} otherwise
    array<pair<int,int>, N> pk;
};
// }}}

#line 1 "Math/multiplicative_function.h"
// NOTE: calculate upto N-1
//
// Multiplicative function {{{
template<int N>
struct MultiplicativeFunction {
    // Init sieve and pk
    MultiplicativeFunction() {
        // Init sieve
        for (int i = 2; i*i < N; i++) {
            if (!sieve[i]) {
                for (int j = i*i; j < N; j += i) {
                    sieve[j] = i;
                }
            }
        }

        // Init pk
        for (int i = 2; i < N; i++) {
            if (!sieve[i]) {
                pk[i] = {i, 1};
            } else {
                int p = sieve[i];

                if (pk[i/p].first == p) { // i = p^k
                    pk[i] = {p, pk[i/p].second + 1};
                } else {
                    pk[i] = {-1, 0};
                }
            }
        }
    }

    // Tested: https://cses.fi/problemset/task/1713
    array<int, N> divisors() {
        array<int, N> res;
        res[1] = 1;

        for (int i = 2; i < N; i++) {
            if (pk[i].first > 0) {  // i = p^k
                res[i] = pk[i].second + 1;
            } else {
                // i = u * v, gcd(u, v) = 1
                int u = i, v = 1;
                int p = sieve[i];
                while (u % p == 0) {
                    u /= p;
                    v *= p;
                }
                res[i] = res[u] * res[v];
            }
        }

        return res;
    }

    // mobius(n) = 1  if n is square-free and has *even* number of prime factors
    // mobius(n) = -1 if n is square-free and has *odd* number of of prime factors
    // mobius(n) = 0  if n is not square-free
    array<int, N> mobius() {
        array<int, N> res;
        res[1] = 1;

        for (int i = 2; i < N; ++i) {
            if (pk[i].first > 0) {  // i = p^k
                res[i] = (pk[i].second >= 2) ? 0 : -1;
            } else {
                // i = u * v, gcd(u, v) = 1
                int u = i, v = 1;
                int p = sieve[i];
                while (u % p == 0) {
                    u /= p;
                    v *= p;
                }
                res[i] = res[u] * res[v];
            }
        }
        return res;
    }

// private:
    // sieve[i] == 0 if i is prime,
    // sieve[i] = any prime factor p otherwise
    array<int, N> sieve = {0};

    // pk[i] = {p, k} if i == p^k
    // pk[i] = {-1, 0} otherwise
    array<pair<int,int>, N> pk;
};
// }}}

Back to top page