FPS Power Projection
(Src/FPS/PowerProjection.hpp)

• View this file on GitHub
• Last update: 2026-01-12 16:27:35+09:00
• Include: #include "Src/FPS/PowerProjection.hpp"

概要

長さ $n$ の数列 $W = (W_0, W_1, \dots, W_{n - 1})$ 、正整数 $m$ 、 FPS $F(x)$ に対して

\[\sum_{i = 0}^{n - 1} w_{i}([x^i]F(x)^{j})\]

を $j = 0, 1, 2, \dots, m - 1$ に対して列挙する。

ライブラリの使い方

template <concepts::IndexedFPS FPS, class Conv = FPSMult>
requires concepts::Convolution<FPS, Conv>
std::vector<typename FPS::value_type> PowerProjection(usize m, FPS W, FPS F, Conv convolution = {}) {

FPS=FPSNTTFriendly<MOD>の場合、convolutionは自動的にoperator*が使われる。それ以外の場合は自分でconvolutionを指定する必要がある。

• $2$ 乗かけて愚直にやる畳み込みはFPS.hppのNaiveConvolutionを使うと良い。

計算量

• $N$ を bit_ceil<usize>(n) とする。
• 長さ $2N$ の列同士の畳み込みを $\log_{2}N$ 回行う
• 長さ $m$ の列同士の畳み込みを高々 $1$ 回行う
• $+O(N\log N + m)$
• つまるところ、FPSNTTFriendly<MOD>では $O(n\log^2 n + m\log m)$

Depends on

• Src/FPS/FPS.hpp
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Verified with

• Test/AtCoder/abc345_g.test.cpp
• Test/AtCoder/abc439_g.test.cpp
• Test/My/FPS/PowerProjection.test.cpp

Code

#pragma once

#include "FPS.hpp"

#include <algorithm>
#include <cassert>
#include <bit>
#include <vector>
#include <ranges>

namespace zawa {

// i = 0, 1, 2, \dots, M - 1
// res[i] = \sum_{j} W([x^j]F(x)^i)
template <concepts::IndexedFPS FPS, class Conv = FPSMult>
requires concepts::Convolution<FPS, Conv>
std::vector<typename FPS::value_type> PowerProjection(usize m, FPS W, FPS F, Conv convolution = {}) {
    using Element = typename FPS::value_type;
    assert(W.size());
    if (!m)
        return {};
    if (F[0] != Element{0}) {
        std::vector<Element> fact(m, 1), invFact(m);
        for (usize i = 1 ; i < m ; i++)
            fact[i] = fact[i - 1] * Element{i};
        invFact[m - 1] = Element{1} / fact[m - 1];
        for (usize i = m ; --i ; )
            invFact[i - 1] = invFact[i] * Element{i};
        const Element c = F[0];
        F[0] = 0;
        FPS cur{PowerProjection(m, W, F, convolution)};
        FPS op(m);
        Element pw = 1;
        for (usize i = 0 ; i < m ; i++, pw *= c) {
            op[i] = pw * invFact[i];
            cur[i] *= invFact[i];
        }
        auto res = convolution(cur, op);
        res.resize(m);
        for (usize i = 0 ; i < m ; i++)
            res[i] *= fact[i];
        return res;
    }
    usize nx = std::bit_ceil(W.size()), ny = 1;
    W.resize(nx);
    F.resize(nx);
    std::ranges::reverse(W);
    // bostan mori [x^{n-1}](g(x)/(1-yf(x)))
    FPS P(2 * nx * ny), Q(2 * nx * ny);
    for (usize x = 0 ; x < nx ; x++) {
        P[x] = W[x];
        Q[x] = -F[x];
    }
    while (nx > 1) {
        FPS R = Q;
        for (usize i = 1 ; i < R.size() ; i += 2)
            R[i] = -R[i];
        auto PR = convolution(P, R), QR = convolution(Q, R);
        PR.resize(4 * nx * ny);
        QR.resize(4 * nx * ny);
        for (usize i = 0 ; i < P.size() ; i++) {
            PR[i + P.size()] += P[i];
            QR[i + P.size()] += Q[i] + R[i];
        }
        std::ranges::fill(P, Element{0});
        std::ranges::fill(Q, Element{0});
        for (usize y = 0 ; y < (ny << 1) ; y++) {
            for (usize x = 1 ; x < nx ; x += 2)
                P[y * nx + (x >> 1)] = PR[y * (nx << 1) + x];
            for (usize x = 0 ; x < nx ; x += 2)
                Q[y * nx + (x >> 1)] = QR[y * (nx << 1) + x];
        }
        nx >>= 1;
        ny <<= 1;
    }
    std::vector<Element> res(ny);
    for (usize i = 0 ; i < ny ; i++)
        res[i] = P[i << 1];
    std::ranges::reverse(res);
    res.resize(m);
    return res;
}

} // namespace zawa

#line 2 "Src/FPS/PowerProjection.hpp"

#line 2 "Src/FPS/FPS.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/FPS/FPS.hpp"

#include <concepts>

namespace zawa {

namespace concepts {

template <class FPS>
concept IndexedFPS = requires(FPS f, usize i) {
    typename FPS::value_type;
    { f.size() } -> std::convertible_to<usize>;
    { f[i] } -> std::convertible_to<typename FPS::value_type>;
    f.reserve(0);
    f.push_back(f[i]);
};

template <class FPS, class Conv>
concept Convolution = 
    std::regular_invocable<Conv, const FPS&, const FPS&> &&
    std::same_as<std::invoke_result_t<Conv, const FPS&, const FPS&>, FPS>;

} // namespace concepts

struct FPSMult {
    template <class FPS>
    requires requires(const FPS& a, const FPS& b) {
        { a * b } -> std::same_as<FPS>;
    }
    FPS operator()(const FPS& a, const FPS& b) const {
        return a * b;
    }
};

struct NaiveConvolution {
    template <class FPS>
    FPS operator()(const FPS& a, const FPS& b) const {
        if (a.empty())
            return b;
        if (b.empty())
            return a;
        FPS res(a.size() + b.size() - 1);
        for (usize i = 0 ; i < a.size() ; i++)
            for (usize j = 0 ; j < b.size() ; j++)
                res[i + j] += a[i] * b[j];
        return res;
    }
};

} // namespace zawa
#line 4 "Src/FPS/PowerProjection.hpp"

#include <algorithm>
#include <cassert>
#include <bit>
#include <vector>
#include <ranges>

namespace zawa {

// i = 0, 1, 2, \dots, M - 1
// res[i] = \sum_{j} W([x^j]F(x)^i)
template <concepts::IndexedFPS FPS, class Conv = FPSMult>
requires concepts::Convolution<FPS, Conv>
std::vector<typename FPS::value_type> PowerProjection(usize m, FPS W, FPS F, Conv convolution = {}) {
    using Element = typename FPS::value_type;
    assert(W.size());
    if (!m)
        return {};
    if (F[0] != Element{0}) {
        std::vector<Element> fact(m, 1), invFact(m);
        for (usize i = 1 ; i < m ; i++)
            fact[i] = fact[i - 1] * Element{i};
        invFact[m - 1] = Element{1} / fact[m - 1];
        for (usize i = m ; --i ; )
            invFact[i - 1] = invFact[i] * Element{i};
        const Element c = F[0];
        F[0] = 0;
        FPS cur{PowerProjection(m, W, F, convolution)};
        FPS op(m);
        Element pw = 1;
        for (usize i = 0 ; i < m ; i++, pw *= c) {
            op[i] = pw * invFact[i];
            cur[i] *= invFact[i];
        }
        auto res = convolution(cur, op);
        res.resize(m);
        for (usize i = 0 ; i < m ; i++)
            res[i] *= fact[i];
        return res;
    }
    usize nx = std::bit_ceil(W.size()), ny = 1;
    W.resize(nx);
    F.resize(nx);
    std::ranges::reverse(W);
    // bostan mori [x^{n-1}](g(x)/(1-yf(x)))
    FPS P(2 * nx * ny), Q(2 * nx * ny);
    for (usize x = 0 ; x < nx ; x++) {
        P[x] = W[x];
        Q[x] = -F[x];
    }
    while (nx > 1) {
        FPS R = Q;
        for (usize i = 1 ; i < R.size() ; i += 2)
            R[i] = -R[i];
        auto PR = convolution(P, R), QR = convolution(Q, R);
        PR.resize(4 * nx * ny);
        QR.resize(4 * nx * ny);
        for (usize i = 0 ; i < P.size() ; i++) {
            PR[i + P.size()] += P[i];
            QR[i + P.size()] += Q[i] + R[i];
        }
        std::ranges::fill(P, Element{0});
        std::ranges::fill(Q, Element{0});
        for (usize y = 0 ; y < (ny << 1) ; y++) {
            for (usize x = 1 ; x < nx ; x += 2)
                P[y * nx + (x >> 1)] = PR[y * (nx << 1) + x];
            for (usize x = 0 ; x < nx ; x += 2)
                Q[y * nx + (x >> 1)] = QR[y * (nx << 1) + x];
        }
        nx >>= 1;
        ny <<= 1;
    }
    std::vector<Element> res(ny);
    for (usize i = 0 ; i < ny ; i++)
        res[i] = P[i << 1];
    std::ranges::reverse(res);
    res.resize(m);
    return res;
}

} // namespace zawa

Back to top page