#879777 (C++23) No.2582 Random Average^K

提出ソース
結果

問題	No.2582 Random Average^K
コンテスト
ユーザー	遭難者
提出日時	2023-06-11 01:09:49
言語	C++23 (gcc 13.3.0 + boost 1.87.0)
結果	RE
実行時間	-
コード長	9,419 bytes
コンパイル時間	6,201 ms
コンパイル使用メモリ	328,008 KB
実行使用メモリ	6,948 KB
最終ジャッジ日時	2024-09-27 03:48:56
合計ジャッジ時間	8,428 ms
ジャッジサーバーID （参考情報）	judge2 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 3
other	AC * 5 RE * 10
権限があれば一括ダウンロードができます
ソースコード

raw source code
#include <bits/stdc++.h>
#include <atcoder/all>
#define rep(i, n) for (int i = 0; i < n; i++)
#define ALL(a) a.begin(), a.end()
#define ll long long
using namespace std;
constexpr int mod = 998244353;
namespace FastFourierTransform
{
	using real = double;

	struct C
	{
		real x, y;

		C() : x(0), y(0) {}

		C(real x, real y) : x(x), y(y) {}

		inline C operator+(const C &c) const { return C(x + c.x, y + c.y); }

		inline C operator-(const C &c) const { return C(x - c.x, y - c.y); }

		inline C operator*(const C &c) const { return C(x * c.x - y * c.y, x * c.y + y * c.x); }

		inline C conj() const { return C(x, -y); }
	};

	const real PI = acosl(-1);
	int base = 1;
	vector<C> rts = {{0, 0}, {1, 0}};
	vector<int> rev = {0, 1};

	void ensure_base(int nbase)
	{
		if (nbase <= base)
			return;
		rev.resize(1 << nbase);
		rts.resize(1 << nbase);
		for (int i = 0; i < (1 << nbase); i++)
		{
			rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (nbase - 1));
		}
		while (base < nbase)
		{
			real angle = PI * 2.0 / (1 << (base + 1));
			for (int i = 1 << (base - 1); i < (1 << base); i++)
			{
				rts[i << 1] = rts[i];
				real angle_i = angle * (2 * i + 1 - (1 << base));
				rts[(i << 1) + 1] = C(cos(angle_i), sin(angle_i));
			}
			++base;
		}
	}

	void fft(vector<C> &a, int n)
	{
		assert((n & (n - 1)) == 0);
		int zeros = __builtin_ctz(n);
		ensure_base(zeros);
		int shift = base - zeros;
		for (int i = 0; i < n; i++)
		{
			if (i < (rev[i] >> shift))
			{
				swap(a[i], a[rev[i] >> shift]);
			}
		}
		for (int k = 1; k < n; k <<= 1)
		{
			for (int i = 0; i < n; i += 2 * k)
			{
				for (int j = 0; j < k; j++)
				{
					C z = a[i + j + k] * rts[j + k];
					a[i + j + k] = a[i + j] - z;
					a[i + j] = a[i + j] + z;
				}
			}
		}
	}

	vector<int64_t> multiply(const vector<int> &a, const vector<int> &b)
	{
		int need = (int)a.size() + (int)b.size() - 1;
		int nbase = 1;
		while ((1 << nbase) < need)
			nbase++;
		ensure_base(nbase);
		int sz = 1 << nbase;
		vector<C> fa(sz);
		for (int i = 0; i < sz; i++)
		{
			int x = (i < (int)a.size() ? a[i] : 0);
			int y = (i < (int)b.size() ? b[i] : 0);
			fa[i] = C(x, y);
		}
		fft(fa, sz);
		C r(0, -0.25 / (sz >> 1)), s(0, 1), t(0.5, 0);
		for (int i = 0; i <= (sz >> 1); i++)
		{
			int j = (sz - i) & (sz - 1);
			C z = (fa[j] * fa[j] - (fa[i] * fa[i]).conj()) * r;
			fa[j] = (fa[i] * fa[i] - (fa[j] * fa[j]).conj()) * r;
			fa[i] = z;
		}
		for (int i = 0; i < (sz >> 1); i++)
		{
			C A0 = (fa[i] + fa[i + (sz >> 1)]) * t;
			C A1 = (fa[i] - fa[i + (sz >> 1)]) * t * rts[(sz >> 1) + i];
			fa[i] = A0 + A1 * s;
		}
		fft(fa, sz >> 1);
		vector<int64_t> ret(need);
		for (int i = 0; i < need; i++)
		{
			ret[i] = llround(i & 1 ? fa[i >> 1].y : fa[i >> 1].x);
		}
		return ret;
	}

};

template <int mod>
struct ModInt
{
	int x;

	ModInt() : x(0) {}

	ModInt(int64_t y) : x(y >= 0 ? y % mod : (mod - (-y) % mod) % mod) {}

	ModInt &operator+=(const ModInt &p)
	{
		if ((x += p.x) >= mod)
			x -= mod;
		return *this;
	}

	ModInt &operator-=(const ModInt &p)
	{
		if ((x += mod - p.x) >= mod)
			x -= mod;
		return *this;
	}

	ModInt &operator*=(const ModInt &p)
	{
		x = (int)(1LL * x * p.x % mod);
		return *this;
	}

	ModInt &operator/=(const ModInt &p)
	{
		*this *= p.inverse();
		return *this;
	}

	ModInt operator-() const { return ModInt(-x); }

	ModInt operator+(const ModInt &p) const { return ModInt(*this) += p; }

	ModInt operator-(const ModInt &p) const { return ModInt(*this) -= p; }

	ModInt operator*(const ModInt &p) const { return ModInt(*this) *= p; }

	ModInt operator/(const ModInt &p) const { return ModInt(*this) /= p; }

	bool operator==(const ModInt &p) const { return x == p.x; }

	bool operator!=(const ModInt &p) const { return x != p.x; }

	ModInt inverse() const
	{
		int a = x, b = mod, u = 1, v = 0, t;
		while (b > 0)
		{
			t = a / b;
			swap(a -= t * b, b);
			swap(u -= t * v, v);
		}
		return ModInt(u);
	}

	ModInt pow(int64_t n) const
	{
		ModInt ret(1), mul(x);
		while (n > 0)
		{
			if (n & 1)
				ret *= mul;
			mul *= mul;
			n >>= 1;
		}
		return ret;
	}

	friend ostream &operator<<(ostream &os, const ModInt &p)
	{
		return os << p.x;
	}

	friend istream &operator>>(istream &is, ModInt &a)
	{
		int64_t t;
		is >> t;
		a = ModInt<mod>(t);
		return (is);
	}

	static int get_mod() { return mod; }
};

using modint = ModInt<mod>;

template <typename T>
struct ArbitraryModConvolution
{
	using real = FastFourierTransform::real;
	using C = FastFourierTransform::C;

	ArbitraryModConvolution() = default;

	vector<T> multiply(const vector<T> &a, const vector<T> &b, int need = -1)
	{
		if (need == -1)
			need = a.size() + b.size() - 1;
		int nbase = 0;
		while ((1 << nbase) < need)
			nbase++;
		FastFourierTransform::ensure_base(nbase);
		int sz = 1 << nbase;
		vector<C> fa(sz);
		for (int i = 0; i < a.size(); i++)
		{
			fa[i] = C(a[i].x & ((1 << 15) - 1), a[i].x >> 15);
		}
		fft(fa, sz);
		vector<C> fb(sz);
		for (int i = 0; i < b.size(); i++)
		{
			fb[i] = C(b[i].x & ((1 << 15) - 1), b[i].x >> 15);
		}
		fft(fb, sz);
		real ratio = 0.25 / sz;
		C r2(0, -1), r3(ratio, 0), r4(0, -ratio), r5(0, 1);
		for (int i = 0; i <= (sz >> 1); i++)
		{
			int j = (sz - i) & (sz - 1);
			C a1 = (fa[i] + fa[j].conj());
			C a2 = (fa[i] - fa[j].conj()) * r2;
			C b1 = (fb[i] + fb[j].conj()) * r3;
			C b2 = (fb[i] - fb[j].conj()) * r4;
			if (i != j)
			{
				C c1 = (fa[j] + fa[i].conj());
				C c2 = (fa[j] - fa[i].conj()) * r2;
				C d1 = (fb[j] + fb[i].conj()) * r3;
				C d2 = (fb[j] - fb[i].conj()) * r4;
				fa[i] = c1 * d1 + c2 * d2 * r5;
				fb[i] = c1 * d2 + c2 * d1;
			}
			fa[j] = a1 * b1 + a2 * b2 * r5;
			fb[j] = a1 * b2 + a2 * b1;
		}
		fft(fa, sz);
		fft(fb, sz);
		vector<T> ret(need);
		for (int i = 0; i < need; i++)
		{
			int64_t aa = llround(fa[i].x);
			int64_t bb = llround(fb[i].x);
			int64_t cc = llround(fa[i].y);
			aa = T(aa).x, bb = T(bb).x, cc = T(cc).x;
			ret[i] = aa + (bb << 15) + (cc << 30);
		}
		return ret;
	}
};

template <typename T>
struct Combination
{
	vector<T> _fact, _rfact, _inv;

	Combination(int sz) : _fact(sz + 1), _rfact(sz + 1), _inv(sz + 1)
	{
		_fact[0] = _rfact[sz] = _inv[0] = 1;
		for (int i = 1; i <= sz; i++)
			_fact[i] = _fact[i - 1] * i;
		_rfact[sz] /= _fact[sz];
		for (int i = sz - 1; i >= 0; i--)
			_rfact[i] = _rfact[i + 1] * (i + 1);
		for (int i = 1; i <= sz; i++)
			_inv[i] = _rfact[i] * _fact[i - 1];
	}

	inline T fact(int k) const { return _fact[k]; }

	inline T rfact(int k) const { return _rfact[k]; }

	inline T inv(int k) const { return _inv[k]; }

	T P(int n, int r) const
	{
		if (r < 0 || n < r)
			return 0;
		return fact(n) * rfact(n - r);
	}

	T C(int p, int q) const
	{
		if (q < 0 || p < q)
			return 0;
		return fact(p) * rfact(q) * rfact(p - q);
	}

	T H(int n, int r) const
	{
		if (n < 0 || r < 0)
			return (0);
		return r == 0 ? 1 : C(n + r - 1, r);
	}
};

template <typename T>
T factorial(int64_t n)
{
	if (n >= mod)
		return 0;
	if (n == 0)
		return 1;

	const int sn = sqrt(n);
	const T sn_inv = T(1) / sn;

	Combination<modint> comb(sn);
	using P = vector<T>;
	ArbitraryModConvolution<modint> fft;

	auto shift = [&](const P &f, T dx)
	{
		int n = (int)f.size();
		T a = dx * sn_inv;
		auto p1 = P(f);
		for (int i = 0; i < n; i++)
		{
			T d = comb.rfact(i) * comb.rfact((n - 1) - i);
			if (((n - 1 - i) & 1))
				d = -d;
			p1[i] *= d;
		}
		auto p2 = P(2 * n);
		for (int i = 0; i < p2.size(); i++)
		{
			p2[i] = (a.x + i - n) <= 0 ? 1 : a + i - n;
		}
		for (int i = 1; i < p2.size(); i++)
		{
			p2[i] *= p2[i - 1];
		}
		T prod = p2[2 * n - 1];
		T prod_inv = T(1) / prod;
		for (int i = 2 * n - 1; i > 0; --i)
		{
			p2[i] = prod_inv * p2[i - 1];
			prod_inv *= a + i - n;
		}
		p2[0] = prod_inv;
		auto p3 = fft.multiply(p1, p2, (int)p2.size());
		p1 = P(p3.begin() + p1.size(), p3.begin() + p2.size());
		prod = 1;
		for (int i = 0; i < n; i++)
		{
			prod *= a + n - 1 - i;
		}
		for (int i = n - 1; i >= 0; --i)
		{
			p1[i] *= prod;
			prod *= p2[n + i] * (a + i - n);
		}
		return p1;
	};
	function<P(int)> rec = [&](int64_t n)
	{
		if (n == 1)
			return P({1, 1 + sn});
		int nh = n >> 1;
		auto a1 = rec(nh);
		auto a2 = shift(a1, nh);
		auto b1 = shift(a1, sn * nh);
		auto b2 = shift(a1, sn * nh + nh);
		for (int i = 0; i <= nh; i++)
			a1[i] *= a2[i];
		for (int i = 1; i <= nh; i++)
			a1.emplace_back(b1[i] * b2[i]);
		if (n & 1)
		{
			for (int64_t i = 0; i < n; i++)
			{
				a1[i] *= n + 1LL * sn * i;
			}
			T prod = 1;
			for (int64_t i = 1LL * n * sn; i < 1LL * n * sn + n; i++)
			{
				prod *= (i + 1);
			}
			a1.push_back(prod);
		}
		return a1;
	};
	auto vs = rec(sn);
	T ret = 1;
	for (int64_t i = 0; i < sn; i++)
		ret *= vs[i];
	for (int64_t i = 1LL * sn * sn + 1; i <= n; i++)
		ret *= i;
	return ret;
}
modint fac[5050], inv[5050];
void solve()
{
	fac[0] = 1;
	for (int i = 1; i < 5050; i++)
		fac[i] = modint(i) * fac[i - 1];
	inv[5049] = fac[5049].inverse();
	for (int i = 5049; i > 0; i--)
		inv[i - 1] = modint(i) * inv[i];
	int n, m;
	cin >> n >> m;
	modint ans = 0;
	for (int k = 1; k <= n; k++)
	{
		modint mul = fac[n] * inv[k] * inv[n - k];
		if ((n - k) & 1)
			mul = -mul;
		mul *= modint(k).pow(n + m);
		ans += mul;
	}
	ans *= factorial<modint>(m);
	ans /= factorial<modint>(n + m) * modint(n).pow(m);
	cout << ans << '\n';
}
int main()
{
	cin.tie(nullptr);
	ios::sync_with_stdio(false);
	cout << fixed << setprecision(13);
	solve();
	return 0;
}
yukicoder

結果

ソースコード
