sps_basic.hpp

• View this file on GitHub
• Last update: 2025-01-19 15:28:01+08:00
• Include: #include "sps_basic.hpp"

Depends on

• subset_conv.hpp

Verified with

• test/set_power_series/exp_of_set_power_series.0.test.cpp
• test/set_power_series/inv_of_set_power_series.0.test.cpp
• test/set_power_series/log_of_set_power_series.0.test.cpp

Code

#pragma once

#include "subset_conv.hpp"
#include <cassert>
#include <vector>

namespace detail {

template<typename Tp> void sps_hadamard_inplace(std::vector<std::vector<Tp>> &rankedA,
                                                const std::vector<std::vector<Tp>> &rankedB,
                                                int LogN) {
    const int N = (1 << LogN);
    for (int i = LogN; i >= 0; --i) {
        for (int j = 0; j < N; ++j) rankedA[i][j] *= rankedB[0][j];
        for (int j = 1; j <= i; ++j)
            for (int k = (1 << j) - 1; k < N; ++k)
                rankedA[i][k] += rankedA[i - j][k] * rankedB[j][k];
    }
}

} // namespace detail

template<typename Tp> inline std::vector<Tp> sps_inv(const std::vector<Tp> &A) {
    const int N = A.size();
    assert(N > 0);
    assert((N & (N - 1)) == 0);
    assert(A[0] != 0);
    int LogN = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector<Tp> res(N);
    res[0] = A[0].inv();
    std::vector rankedInvA(LogN, std::vector<Tp>(N / 2));
    for (int i = 0; i < LogN; ++i) {
        std::vector rankedA(i + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j < (1 << i); ++j) rankedA[__builtin_popcount(j)][j] = A[j + (1 << i)];
        for (int j = 0; j <= i; ++j) subset_zeta(rankedA[j]);

        for (int j = (1 << i) / 2; j < (1 << i); ++j) rankedInvA[__builtin_popcount(j)][j] = res[j];
        for (int j = 0; j <= i; ++j) {
            subset_zeta_n(rankedInvA[j].begin() + (1 << i) / 2, (1 << i) / 2);
            for (int k = 0; k < (1 << i) / 2; ++k)
                rankedInvA[j][k + (1 << i) / 2] += rankedInvA[j][k];
        }

        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);
        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);

        for (int j = 0; j <= i; ++j) subset_moebius(rankedA[j]);
        for (int j = 0; j < (1 << i); ++j) res[j + (1 << i)] = -rankedA[__builtin_popcount(j)][j];
    }
    return res;
}

template<typename Tp> inline std::vector<Tp> sps_log(const std::vector<Tp> &A) {
    const int N = A.size();
    assert(N > 0);
    assert((N & (N - 1)) == 0);
    assert(A[0] == 1);
    int LogN = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector<Tp> res(N);
    std::vector<Tp> invA = {Tp(1)};
    invA.resize(N / 2);
    std::vector rankedInvA(LogN, std::vector<Tp>(N / 2));
    for (int i = 0; i < LogN; ++i) {
        std::vector rankedA(i + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j < (1 << i); ++j) rankedA[__builtin_popcount(j)][j] = A[j + (1 << i)];
        for (int j = 0; j <= i; ++j) subset_zeta(rankedA[j]);

        for (int j = (1 << i) / 2; j < (1 << i); ++j)
            rankedInvA[__builtin_popcount(j)][j] = invA[j];
        for (int j = 0; j <= i; ++j) {
            subset_zeta_n(rankedInvA[j].begin() + (1 << i) / 2, (1 << i) / 2);
            for (int k = 0; k < (1 << i) / 2; ++k)
                rankedInvA[j][k + (1 << i) / 2] += rankedInvA[j][k];
        }

        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);
        auto rankedLogA = rankedA;
        for (int j = 0; j <= i; ++j) subset_moebius(rankedLogA[j]);
        for (int j = 0; j < (1 << i); ++j) res[j + (1 << i)] = rankedLogA[__builtin_popcount(j)][j];
        if (i == LogN - 1) break;
        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);

        for (int j = 0; j <= i; ++j) subset_moebius(rankedA[j]);
        for (int j = 0; j < (1 << i); ++j) invA[j + (1 << i)] = -rankedA[__builtin_popcount(j)][j];
    }
    return res;
}

// returns exp(0 + tx_1 + ...) in R[x_1,...,x_n]/(x_1^2,...,x_n^2)
template<typename Tp> inline std::vector<Tp> sps_exp(const std::vector<Tp> &A) {
    const int N = A.size();
    assert(N > 0);
    assert((N & (N - 1)) == 0);
    assert(A[0] == 0);
    int LogN = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector<Tp> res(N);
    res[0] = 1;
    std::vector rankedExpA(LogN, std::vector<Tp>(N / 2));
    for (int i = 0; i < LogN; ++i) {
        std::vector rankedA(i + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j < (1 << i); ++j) rankedA[__builtin_popcount(j)][j] = A[j + (1 << i)];
        for (int j = 0; j <= i; ++j) subset_zeta(rankedA[j]);

        for (int j = (1 << i) / 2; j < (1 << i); ++j) rankedExpA[__builtin_popcount(j)][j] = res[j];
        for (int j = 0; j <= i; ++j) {
            subset_zeta_n(rankedExpA[j].begin() + (1 << i) / 2, (1 << i) / 2);
            for (int k = 0; k < (1 << i) / 2; ++k)
                rankedExpA[j][k + (1 << i) / 2] += rankedExpA[j][k];
        }

        std::vector<int> map(i + 1);
        for (int j = 0; j <= i; ++j) map[j] = (j & 1) ? map[j / 2] : j / 2;

        std::vector ExpAA(i / 2 + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j <= i; ++j)
            for (int k = 0; j + k <= i; ++k)
                for (int l = (1 << k) - 1; l < (1 << i); ++l)
                    ExpAA[map[j + k]][l] += rankedExpA[j][l] * rankedA[k][l];
        for (int j = 0; j <= i / 2; ++j) subset_moebius(ExpAA[j]);

        for (int j = 0; j < (1 << i); ++j) res[j + (1 << i)] = ExpAA[map[__builtin_popcount(j)]][j];
    }
    return res;
}

#line 2 "sps_basic.hpp"

#line 2 "subset_conv.hpp"

#include <cassert>
#include <vector>

template<typename Tp> inline std::vector<std::vector<Tp>> to_ranked(const std::vector<Tp> &A) {
    const int N = A.size();
    int LogN    = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector res(LogN + 1, std::vector<Tp>(N));
    for (int i = 0; i < N; ++i) res[__builtin_popcount(i)][i] = A[i];
    return res;
}

template<typename Tp> inline std::vector<Tp> from_ranked(const std::vector<std::vector<Tp>> &A) {
    const int N = A[0].size();
    std::vector<Tp> res(N);
    for (int i = 0; i < N; ++i) res[i] = A[__builtin_popcount(i)][i];
    return res;
}

template<typename Iterator> inline void subset_zeta_n(Iterator a, int n) {
    assert((n & (n - 1)) == 0);
    for (int i = 2; i <= n; i *= 2)
        for (int j = 0; j < n; j += i)
            for (int k = j; k < j + i / 2; ++k) a[k + i / 2] += a[k];
}

template<typename Tp> inline void subset_zeta(std::vector<Tp> &a) {
    subset_zeta_n(a.begin(), a.size());
}

template<typename Iterator> inline void subset_moebius_n(Iterator a, int n) {
    assert((n & (n - 1)) == 0);
    for (int i = 2; i <= n; i *= 2)
        for (int j = 0; j < n; j += i)
            for (int k = j; k < j + i / 2; ++k) a[k + i / 2] -= a[k];
}

template<typename Tp> inline void subset_moebius(std::vector<Tp> &a) {
    subset_moebius_n(a.begin(), a.size());
}

template<typename Tp>
inline std::vector<Tp> subset_convolution(const std::vector<Tp> &A, const std::vector<Tp> &B) {
    assert(A.size() == B.size());
    const int N = A.size();
    int LogN    = 0;
    while ((1 << LogN) != N) ++LogN;
    auto rankedA = to_ranked(A);
    auto rankedB = to_ranked(B);

    for (int i = 0; i <= LogN; ++i) {
        subset_zeta(rankedA[i]);
        subset_zeta(rankedB[i]);
    }

    // see: https://codeforces.com/blog/entry/126418
    // see: https://oeis.org/A025480
    std::vector<int> map(LogN + 1);
    for (int i = 0; i <= LogN; ++i) map[i] = (i & 1) ? map[i / 2] : i / 2;

    std::vector rankedAB(LogN / 2 + 1, std::vector<Tp>(N));
    for (int i = 0; i <= LogN; ++i)
        for (int j = 0; i + j <= LogN; ++j)
            for (int k = (1 << j) - 1; k < N; ++k)
                rankedAB[map[i + j]][k] += rankedA[i][k] * rankedB[j][k];

    for (int i = 0; i <= LogN / 2; ++i) subset_moebius(rankedAB[i]);

    std::vector<Tp> res(N);
    for (int i = 0; i < N; ++i) res[i] = rankedAB[map[__builtin_popcount(i)]][i];
    return res;
}
#line 6 "sps_basic.hpp"

namespace detail {

template<typename Tp> void sps_hadamard_inplace(std::vector<std::vector<Tp>> &rankedA,
                                                const std::vector<std::vector<Tp>> &rankedB,
                                                int LogN) {
    const int N = (1 << LogN);
    for (int i = LogN; i >= 0; --i) {
        for (int j = 0; j < N; ++j) rankedA[i][j] *= rankedB[0][j];
        for (int j = 1; j <= i; ++j)
            for (int k = (1 << j) - 1; k < N; ++k)
                rankedA[i][k] += rankedA[i - j][k] * rankedB[j][k];
    }
}

} // namespace detail

template<typename Tp> inline std::vector<Tp> sps_inv(const std::vector<Tp> &A) {
    const int N = A.size();
    assert(N > 0);
    assert((N & (N - 1)) == 0);
    assert(A[0] != 0);
    int LogN = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector<Tp> res(N);
    res[0] = A[0].inv();
    std::vector rankedInvA(LogN, std::vector<Tp>(N / 2));
    for (int i = 0; i < LogN; ++i) {
        std::vector rankedA(i + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j < (1 << i); ++j) rankedA[__builtin_popcount(j)][j] = A[j + (1 << i)];
        for (int j = 0; j <= i; ++j) subset_zeta(rankedA[j]);

        for (int j = (1 << i) / 2; j < (1 << i); ++j) rankedInvA[__builtin_popcount(j)][j] = res[j];
        for (int j = 0; j <= i; ++j) {
            subset_zeta_n(rankedInvA[j].begin() + (1 << i) / 2, (1 << i) / 2);
            for (int k = 0; k < (1 << i) / 2; ++k)
                rankedInvA[j][k + (1 << i) / 2] += rankedInvA[j][k];
        }

        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);
        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);

        for (int j = 0; j <= i; ++j) subset_moebius(rankedA[j]);
        for (int j = 0; j < (1 << i); ++j) res[j + (1 << i)] = -rankedA[__builtin_popcount(j)][j];
    }
    return res;
}

template<typename Tp> inline std::vector<Tp> sps_log(const std::vector<Tp> &A) {
    const int N = A.size();
    assert(N > 0);
    assert((N & (N - 1)) == 0);
    assert(A[0] == 1);
    int LogN = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector<Tp> res(N);
    std::vector<Tp> invA = {Tp(1)};
    invA.resize(N / 2);
    std::vector rankedInvA(LogN, std::vector<Tp>(N / 2));
    for (int i = 0; i < LogN; ++i) {
        std::vector rankedA(i + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j < (1 << i); ++j) rankedA[__builtin_popcount(j)][j] = A[j + (1 << i)];
        for (int j = 0; j <= i; ++j) subset_zeta(rankedA[j]);

        for (int j = (1 << i) / 2; j < (1 << i); ++j)
            rankedInvA[__builtin_popcount(j)][j] = invA[j];
        for (int j = 0; j <= i; ++j) {
            subset_zeta_n(rankedInvA[j].begin() + (1 << i) / 2, (1 << i) / 2);
            for (int k = 0; k < (1 << i) / 2; ++k)
                rankedInvA[j][k + (1 << i) / 2] += rankedInvA[j][k];
        }

        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);
        auto rankedLogA = rankedA;
        for (int j = 0; j <= i; ++j) subset_moebius(rankedLogA[j]);
        for (int j = 0; j < (1 << i); ++j) res[j + (1 << i)] = rankedLogA[__builtin_popcount(j)][j];
        if (i == LogN - 1) break;
        detail::sps_hadamard_inplace(rankedA, rankedInvA, i);

        for (int j = 0; j <= i; ++j) subset_moebius(rankedA[j]);
        for (int j = 0; j < (1 << i); ++j) invA[j + (1 << i)] = -rankedA[__builtin_popcount(j)][j];
    }
    return res;
}

// returns exp(0 + tx_1 + ...) in R[x_1,...,x_n]/(x_1^2,...,x_n^2)
template<typename Tp> inline std::vector<Tp> sps_exp(const std::vector<Tp> &A) {
    const int N = A.size();
    assert(N > 0);
    assert((N & (N - 1)) == 0);
    assert(A[0] == 0);
    int LogN = 0;
    while ((1 << LogN) != N) ++LogN;
    std::vector<Tp> res(N);
    res[0] = 1;
    std::vector rankedExpA(LogN, std::vector<Tp>(N / 2));
    for (int i = 0; i < LogN; ++i) {
        std::vector rankedA(i + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j < (1 << i); ++j) rankedA[__builtin_popcount(j)][j] = A[j + (1 << i)];
        for (int j = 0; j <= i; ++j) subset_zeta(rankedA[j]);

        for (int j = (1 << i) / 2; j < (1 << i); ++j) rankedExpA[__builtin_popcount(j)][j] = res[j];
        for (int j = 0; j <= i; ++j) {
            subset_zeta_n(rankedExpA[j].begin() + (1 << i) / 2, (1 << i) / 2);
            for (int k = 0; k < (1 << i) / 2; ++k)
                rankedExpA[j][k + (1 << i) / 2] += rankedExpA[j][k];
        }

        std::vector<int> map(i + 1);
        for (int j = 0; j <= i; ++j) map[j] = (j & 1) ? map[j / 2] : j / 2;

        std::vector ExpAA(i / 2 + 1, std::vector<Tp>(1 << i));
        for (int j = 0; j <= i; ++j)
            for (int k = 0; j + k <= i; ++k)
                for (int l = (1 << k) - 1; l < (1 << i); ++l)
                    ExpAA[map[j + k]][l] += rankedExpA[j][l] * rankedA[k][l];
        for (int j = 0; j <= i / 2; ++j) subset_moebius(ExpAA[j]);

        for (int j = 0; j < (1 << i); ++j) res[j + (1 << i)] = ExpAA[map[__builtin_popcount(j)]][j];
    }
    return res;
}

Back to top page