// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// Partition.hpp
//
// 分割関数 (Partition Function) と関連関数
//
// p(n): n の分割数 (n を正整数の和として表す方法の数)
//
// 主な API:
//   partition(n)              — p(n) (DP)
//   partitionTable(n)         — p(0)..p(n) のテーブル
//   partitionQ(n)             — q(n) 重複なし分割数
//   pentagonalNumber(k)       — 一般五角数
//   partitions(n)             — 分割の列挙

#ifndef CALX_CORE_PARTITION_HPP
#define CALX_CORE_PARTITION_HPP

#include <cstdint>
#include <vector>
#include <algorithm>
#include <stdexcept>

namespace calx {

/**
 * @brief 五角数 ω(k) = k(3k-1)/2
 *
 * 一般五角数: k = 0, 1, -1, 2, -2, ...
 * → ω(k) = 0, 1, 2, 5, 7, 12, 15, 22, 26, ...
 */
[[nodiscard]]
constexpr int64_t pentagonalNumber(int64_t k) {
    return k * (3 * k - 1) / 2;
}

/**
 * @brief 分割数 p(n) のテーブル (n = 0, ..., maxN)
 *
 * Euler の五角数定理による O(n√n) アルゴリズム:
 * p(n) = Σ_{k≠0} (-1)^{k+1} · p(n - ω(k))
 */
[[nodiscard]]
inline std::vector<int64_t> partitionTable(int64_t maxN) {
    if (maxN < 0) return {};

    std::vector<int64_t> p(static_cast<std::size_t>(maxN + 1), 0);
    p[0] = 1;

    for (int64_t n = 1; n <= maxN; ++n) {
        int64_t sum = 0;
        for (int64_t k = 1; ; ++k) {
            // 一般五角数: k(3k-1)/2 と k(3k+1)/2 = (-k)(3(-k)-1)/2
            int64_t pent1 = k * (3 * k - 1) / 2;
            int64_t pent2 = k * (3 * k + 1) / 2;

            if (pent1 > n) break;

            int64_t sign = (k % 2 == 1) ? 1 : -1;
            sum += sign * p[static_cast<std::size_t>(n - pent1)];

            if (pent2 <= n) {
                sum += sign * p[static_cast<std::size_t>(n - pent2)];
            }
        }
        p[static_cast<std::size_t>(n)] = sum;
    }

    return p;
}

/**
 * @brief 分割数 p(n)
 *
 * p(0) = 1, p(1) = 1, p(2) = 2, p(3) = 3, p(4) = 5, p(5) = 7, ...
 */
[[nodiscard]]
inline int64_t partition(int64_t n) {
    if (n < 0) return 0;
    if (n == 0) return 1;
    auto table = partitionTable(n);
    return table[static_cast<std::size_t>(n)];
}

/**
 * @brief 重複なし分割数 q(n)
 *
 * n を相異なる正整数の和として表す方法の数。
 * DP: q(n, k) = q(n, k-1) + q(n-k, k-1)
 */
[[nodiscard]]
inline int64_t partitionQ(int64_t n) {
    if (n < 0) return 0;
    if (n == 0) return 1;

    // DP: q[j] = 相異なる部分で和が j となる方法数
    std::vector<int64_t> q(static_cast<std::size_t>(n + 1), 0);
    q[0] = 1;

    for (int64_t k = 1; k <= n; ++k) {
        // 逆順に更新 (0-1 ナップサック)
        for (int64_t j = n; j >= k; --j) {
            q[static_cast<std::size_t>(j)] += q[static_cast<std::size_t>(j - k)];
        }
    }
    return q[static_cast<std::size_t>(n)];
}

/**
 * @brief n の分割を全列挙
 *
 * 各分割は降順のベクトル。n が小さい場合にのみ使用。
 */
[[nodiscard]]
inline std::vector<std::vector<int64_t>> partitions(int64_t n) {
    std::vector<std::vector<int64_t>> result;
    if (n < 0) return result;
    if (n == 0) { result.push_back({}); return result; }

    // 再帰的列挙
    std::vector<int64_t> current;

    auto generate = [&](auto& self, int64_t remaining, int64_t maxPart) -> void {
        if (remaining == 0) {
            result.push_back(current);
            return;
        }
        for (int64_t k = std::min(remaining, maxPart); k >= 1; --k) {
            current.push_back(k);
            self(self, remaining - k, k);
            current.pop_back();
        }
    };

    generate(generate, n, n);
    return result;
}

/**
 * @brief 分割数 p(n, k): ちょうど k 個の部分に分ける分割数
 */
[[nodiscard]]
inline int64_t partitionIntoKParts(int64_t n, int64_t k) {
    if (k <= 0 || n < k) return 0;
    if (k == 1 || n == k) return 1;

    // p(n,k) = p(n-1,k-1) + p(n-k,k)
    // DP テーブル
    std::vector<std::vector<int64_t>> dp(
        static_cast<std::size_t>(n + 1),
        std::vector<int64_t>(static_cast<std::size_t>(k + 1), 0));

    for (int64_t i = 1; i <= n; ++i) dp[static_cast<std::size_t>(i)][1] = 1;
    for (int64_t i = 1; i <= n; ++i) {
        for (int64_t j = 2; j <= std::min(i, k); ++j) {
            dp[static_cast<std::size_t>(i)][static_cast<std::size_t>(j)] =
                dp[static_cast<std::size_t>(i - 1)][static_cast<std::size_t>(j - 1)] +
                dp[static_cast<std::size_t>(i - j)][static_cast<std::size_t>(j)];
        }
    }
    return dp[static_cast<std::size_t>(n)][static_cast<std::size_t>(k)];
}

/**
 * @brief Euler の分割恒等式の検証
 *
 * 奇数部分のみの分割数 = 相異なる部分の分割数
 */
[[nodiscard]]
inline int64_t partitionOddParts(int64_t n) {
    if (n < 0) return 0;
    if (n == 0) return 1;

    std::vector<int64_t> q(static_cast<std::size_t>(n + 1), 0);
    q[0] = 1;

    for (int64_t k = 1; k <= n; k += 2) {
        for (int64_t j = k; j <= n; ++j) {
            q[static_cast<std::size_t>(j)] += q[static_cast<std::size_t>(j - k)];
        }
    }
    return q[static_cast<std::size_t>(n)];
}

} // namespace calx

#endif // CALX_CORE_PARTITION_HPP
