cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:warning: 個数指定の部分和集合の内部処理

components/fixed_size_subset_sum.hpp · GitHub で見る

共通ページ

個数指定の部分和集合の内部部品です。通常は math/fixed_size_subset_sum.hpp を include します。

\(\displaystyle F(x)=\sum_{a\in S}x^a\) の順変換から、\(F(x^j)\) の各周波数での値を添字の置換で取得します。各周波数で Newton の恒等式 \(\displaystyle i e_i=\sum_{j=1}^{i}(-1)^{j-1}e_{i-j}F(x^j)\) を計算し、\(e_K\) だけを逆変換します。\(e_K\) の係数は、異なる要素を \(K\) 個選ぶ方法の個数です。逆変換後の共通の非ゼロ定数倍は、存在判定には影響しません。

一つの和を作る選び方が \(c\) 通りなら、そのそれぞれから一要素を除いた集合は合計 \(Kc\) 個あり、同じ \(K-1\) 要素の集合は二度現れません。残りの一要素が和から一意に決まるためです。したがって \(\displaystyle c\leq\left\lfloor\frac{\binom m{K-1}}K\right\rfloor\) です。

まだ余りが 0 の和だけについて、候補になれる要素の個数から上限を絞ります。和を \(s\)、集合全体で小さい方・大きい方から \(K-1\) 個の和をそれぞれ \(A,B\) とすると、選ばれる各要素は \([s-B,s-A]\) に入ります。この区間内の要素数を先ほどの上限に代入し、使用した素数の積より小さくなれば追加の変換は不要です。

素数 原始根
998244353 3
2130706433 3
2113929217 5
2088763393 5
2013265921 31
1811939329 13
1711276033 29
1484783617 5
1300234241 3
1224736769 3

いずれも長さ \(2^{23}\) の NTT に対応します。

周波数の数を \(L\) とすると、素数一つにつき \(O(L(\log L+K^2))\) 時間、\(O(L+K)\) 空間です。

依存ライブラリ

コード

#ifndef FIXED_SIZE_SUBSET_SUM_CORE_HPP
#define FIXED_SIZE_SUBSET_SUM_CORE_HPP

#include "ntt_mod.hpp"
#include <array>
#include <stdexcept>

namespace fixed_size_subset_sum_detail {
using Word = std::uint32_t;
using Digits = std::vector<Word>;
inline void multiply(Digits& a, Word x) {
	std::uint64_t carry = 0;
	for (auto& v : a) {
		carry += std::uint64_t(v) * x;
		v = Word(carry);
		carry >>= 32;
	}
	if (carry) a.push_back(Word(carry));
}
inline void divide(Digits& a, Word x) {
	std::uint64_t carry = 0;
	for (std::size_t i = a.size(); i--;) {
		carry = (carry << 32) | a[i];
		a[i] = Word(carry / x);
		carry %= x;
	}
	while (a.size() > 1 && !a.back()) a.pop_back();
}
inline bool less(const Digits& a, const Digits& b) {
	if (a.size() != b.size()) return a.size() < b.size();
	return std::lexicographical_compare(a.rbegin(), a.rend(), b.rbegin(), b.rend());
}

template<Word Mod, Word Root>
std::size_t include_modular(const std::vector<int>& values, int k, int low, std::vector<unsigned char>& present) {
	std::size_t size = std::bit_ceil(std::size_t(low) + present.size());
	bool simd = false;
#if defined(__x86_64__) || defined(__i386__)
	simd = __builtin_cpu_supports("avx2");
#endif
	ntt_mod_detail::Plan<Mod, Root> plan(size, simd);
	std::vector<Word> frequency(size), result(size), powers(k + 1), elementary(k + 1), inverse(k + 1, 1);
	for (int x : values) frequency[x] = 1;
	ntt_mod_detail::forward(frequency, plan, simd);
	if constexpr (Mod < (Word(1) << 30)) for (auto& x : frequency) if (x >= Mod) x -= Mod;
	ntt_mod_detail::bit_reverse(frequency);
	for (int i = 2; i <= k; ++i) inverse[i] = Mod - std::uint64_t(Mod / i) * inverse[Mod % i] % Mod;
	elementary[0] = 1;
	for (std::size_t x = 0; x < size; ++x) {
		for (int j = 1; j <= k; ++j) {
			Word value = frequency[(x * j) & (size - 1)];
			powers[j] = (j & 1) || !value ? value : Mod - value;
		}
		elementary[1] = powers[1];
		for (int i = 2; i <= k; ++i) {
			__uint128_t sum = powers[i];
			for (int j = 1; j < i; ++j) sum += std::uint64_t(elementary[i - j]) * powers[j];
			elementary[i] = std::uint64_t(sum % Mod) * inverse[i] % Mod;
		}
		result[x] = elementary[k];
	}
	ntt_mod_detail::bit_reverse(result);
	ntt_mod_detail::backward(result, plan, simd);
	std::reverse(result.begin() + 1, result.end());
	std::size_t added = 0;
	for (std::size_t i = 0; i < present.size(); ++i) {
		Word x = result[low + i];
		bool nonzero = x && x != Mod;
		added += !present[i] && nonzero;
		present[i] |= nonzero;
	}
	return added;
}

struct Pass {
	Word modulus;
	std::size_t (*include)(const std::vector<int>&, int, int, std::vector<unsigned char>&);
};
inline constexpr std::array passes{
	Pass{998244353, include_modular<998244353, 3>},
	Pass{2130706433, include_modular<2130706433, 3>},
	Pass{2113929217, include_modular<2113929217, 5>},
	Pass{2088763393, include_modular<2088763393, 5>},
	Pass{2013265921, include_modular<2013265921, 31>},
	Pass{1811939329, include_modular<1811939329, 13>},
	Pass{1711276033, include_modular<1711276033, 29>},
	Pass{1484783617, include_modular<1484783617, 5>},
	Pass{1300234241, include_modular<1300234241, 3>},
	Pass{1224736769, include_modular<1224736769, 3>}
};

inline Digits coefficient_bound(int n, int k) {
	if (n < k) return {0};
	Digits bound{1}, limit{1};
	for (auto pass : passes) multiply(limit, pass.modulus);
	multiply(limit, Word(k));
	int degree = std::min(k - 1, n - k + 1);
	for (int i = 1; i <= degree; ++i) {
		multiply(bound, Word(n - i + 1));
		divide(bound, Word(i));
		if (!less(bound, limit)) throw std::length_error("too many subsets");
	}
	divide(bound, Word(k));
	return bound;
}
}

#endif
Back to top page