Subset 畳み込み

convolution/subset_convolution.hpp · GitHub で見る

共通ページ

#include "lib/convolution/subset_convolution.hpp"

std::vector<long long> a{1, 2, 3, 4}, b{5, 6, 7, 8};
auto c = subset_convolution(a, b); // {5,16,22,60}

操作	内容・返り値	計算量
subset_convolution(a,b)	\(\displaystyle c_S=\sum_{A\subseteq S}a_Ab_{S\setminus A}\)	\(O(N\log^2 N)\)

添字を bitmask として扱い、重ならない 2 集合に分けた全ての組の積を加えます。例の添字 3 は集合 {0,1} で、c[3] = a[0]*b[3]+a[1]*b[2]+a[2]*b[1]+a[3]*b[0] = 60 です。入力を変更せず、新しい配列を返します。足りない係数は 0 とし、どちらかが空なら空配列です。

\(\displaystyle c_S = \sum_{A\subseteq S} a_A b_{S\setminus A}\)。添字を部分集合の bit 表現として扱います。係数は環の要素とし、長い方の長さを 2 の冪に切り上げた \(N\) 個を返します。

時間 \(O(N \log^{2} N)\)、追加空間 \(O(N \log N)\) です。

依存ライブラリ

• Subset Convolution Modular Core

検証コード

• test/subset_convolution.test.cpp

コード

#ifndef SUBSET_CONVOLUTION_HPP
#define SUBSET_CONVOLUTION_HPP

#include <algorithm>
#include <bit>
#include <cstddef>
#include <vector>
#include "../components/subset_convolution_modular.hpp"

#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif

template <class T>
std::vector<T> subset_convolution(const std::vector<T> &a, const std::vector<T> &b) {
	if (a.empty() || b.empty()) return {};
	std::size_t n = std::bit_ceil(std::max(a.size(), b.size()));
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
	if constexpr (subset_convolution_detail::SmallModint<T>::value)
		if (n >= 256 && n <= (std::size_t(1) << 20) && subset_convolution_detail::avx2())
			return subset_convolution_detail::Modular<T::mod()>::convolution(a, b, int(n));
#endif
	int levels = std::bit_width(n) - 1, width = levels + 1;
	std::vector<T> x(n * width, T(0)), y(n * width, T(0)), result(n, T(0));
	for (std::size_t i = 0; i < a.size(); ++i) x[i * width + std::popcount(i)] = a[i];
	for (std::size_t i = 0; i < b.size(); ++i) y[i * width + std::popcount(i)] = b[i];
	auto zeta = [&](std::size_t first, std::size_t length, std::size_t half) {
		for (; half < length; half *= 2)
			for (std::size_t start = first; start < first + length; start += 2 * half) {
				int low = std::popcount(start);
				for (std::size_t j = start; j < start + half; ++j)
					for (int k = low, end = std::popcount(j); k <= end; ++k) {
						x[(j + half) * width + k] += x[j * width + k];
						y[(j + half) * width + k] += y[j * width + k];
					}
			}
	};
	std::size_t block = std::min<std::size_t>(n, 256);
	for (std::size_t start = 0; start < n; start += block) zeta(start, block, 1);
	zeta(0, n, block);
	for (std::size_t mask = 0; mask < n; ++mask) {
		int count = std::popcount(mask);
		for (int k = std::min(levels, 2 * count); k >= count; --k) {
			T sum(0);
			for (int i = k - count; i <= count; ++i)
				sum += x[mask * width + i] * y[mask * width + k - i];
			x[mask * width + k] = sum;
		}
	}
	auto mobius = [&](std::size_t first, std::size_t length, std::size_t half) {
		for (; half < length; half *= 2)
			for (std::size_t start = first; start < first + length; start += 2 * half) {
				int low = std::popcount(start) + 1, high = levels - std::countr_zero(half);
				for (std::size_t j = start; j < start + half; ++j) {
					int count = std::popcount(j - start);
					for (int k = low + count; k <= high + count; ++k)
						x[(j + half) * width + k] -= x[j * width + k];
				}
			}
	};
	for (std::size_t start = 0; start < n; start += block) mobius(start, block, 1);
	mobius(0, n, block);
	for (std::size_t mask = 0; mask < n; ++mask) result[mask] = x[mask * width + std::popcount(mask)];
	return result;
}

#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif

#endif

Back to top page