結果

問題 No.2005 Sum of Power Sums
ユーザー 👑 ygussanyygussany
提出日時 2022-06-20 23:10:12
言語 C
(gcc 12.3.0)
結果
TLE  
(最新)
AC  
(最初)
実行時間 -
コード長 9,218 bytes
コンパイル時間 1,425 ms
コンパイル使用メモリ 57,568 KB
実行使用メモリ 79,104 KB
最終ジャッジ日時 2024-04-21 13:37:12
合計ジャッジ時間 11,583 ms
ジャッジサーバーID
(参考情報)
judge4 / judge1
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 1 ms
5,376 KB
testcase_02 AC 28 ms
5,376 KB
testcase_03 AC 1 ms
5,376 KB
testcase_04 AC 1 ms
5,376 KB
testcase_05 AC 1 ms
5,376 KB
testcase_06 AC 1 ms
5,376 KB
testcase_07 AC 1 ms
5,376 KB
testcase_08 AC 2 ms
5,376 KB
testcase_09 AC 8 ms
5,376 KB
testcase_10 AC 1 ms
5,376 KB
testcase_11 AC 24 ms
5,376 KB
testcase_12 AC 1 ms
5,376 KB
testcase_13 AC 24 ms
5,376 KB
testcase_14 AC 950 ms
77,184 KB
testcase_15 TLE -
testcase_16 TLE -
testcase_17 TLE -
testcase_18 AC 1 ms
5,376 KB
testcase_19 AC 24 ms
5,376 KB
testcase_20 TLE -
権限があれば一括ダウンロードができます

ソースコード

diff #

#pragma GCC target("avx2")
#pragma GCC optimize("Ofast,unroll-loops")
#include <stdio.h>
#include <stdlib.h>

const int Mod = 998244353,
	bit[21] = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288, 1048576},
	bit_inv[21] = {1, 499122177, 748683265, 873463809, 935854081, 967049217, 982646785, 990445569, 994344961, 996294657, 997269505, 997756929, 998000641, 998122497, 998183425, 998213889, 998229121, 998236737, 998240545, 998242449, 998243401},
	root[21] = {1, 998244352, 911660635, 372528824, 929031873, 452798380, 922799308, 781712469, 476477967, 166035806, 258648936, 584193783, 63912897, 350007156, 666702199, 968855178, 629671588, 24514907, 996173970, 363395222, 565042129},
	root_inv[21] = {1, 998244352, 86583718, 509520358, 337190230, 87557064, 609441965, 135236158, 304459705, 685443576, 381598368, 335559352, 129292727, 358024708, 814576206, 708402881, 283043518, 3707709, 121392023, 704923114, 950391366};
int ntt_b[20][524289], ntt_c[20][524289], ntt_x[20][524289], ntt_y[20][524289];

long long div_mod(long long x, long long y, long long z)
{
	if (x % y == 0) return x / y;
	else return (div_mod((1 + x / y) * y - x, (z % y), y) * z + x) / y;
}

void NTT(int k, int a[], int z[])
{
	if (k == 0) {
		z[0] = a[0];
		return;
	}
	
	int i, d = bit[k-1], tmpp;
	long long tmp;
	for (i = 0; i < d; i++) {
		ntt_b[k][i] = a[i*2];
		ntt_c[k][i] = a[i*2+1];
	}
	NTT(k - 1, ntt_b[k], ntt_x[k]);
	NTT(k - 1, ntt_c[k], ntt_y[k]);
	for (i = 0, tmp = 1; i < d; i++, tmp = tmp * root[k] % Mod) {
		tmpp = tmp * ntt_y[k][i] % Mod;
		z[i] = ntt_x[k][i] + tmpp;
		if (z[i] >= Mod) z[i] -= Mod;
		z[i+d] = ntt_x[k][i] - tmpp;
		if (z[i+d] < 0) z[i+d] += Mod;
	}
}

void NTT_reverse(int k, int z[], int a[])
{
	if (k == 0) {
		a[0] = z[0];
		return;
	}
	
	int i, d = bit[k-1], tmpp;
	long long tmp;
	for (i = 0; i < d; i++) {
		ntt_x[k][i] = z[i*2];
		ntt_y[k][i] = z[i*2+1];
	}
	NTT_reverse(k - 1, ntt_x[k], ntt_b[k]);
	NTT_reverse(k - 1, ntt_y[k], ntt_c[k]);
	for (i = 0, tmp = 1; i < d; i++, tmp = tmp * root_inv[k] % Mod) {
		tmpp = tmp * ntt_c[k][i] % Mod;
		a[i] = ntt_b[k][i] + tmpp;
		if (a[i] >= Mod) a[i] -= Mod;
		a[i+d] = ntt_b[k][i] - tmpp;
		if (a[i+d] < 0) a[i+d] += Mod;
	}
}
// Compute the product of two polynomials a[0-da] and b[0-db] using NTT in O(d * log d) time
void prod_poly_NTT(int da, int db, int a[], int b[], int c[])
{
	int i, k;
	static int aa[524289], bb[524289], cc[524289];
	for (k = 0; bit[k] <= da + db; k++);
	for (i = 0; i <= da; i++) aa[i] = a[i];
	for (i = da + 1; i < bit[k]; i++) aa[i] = 0;
	for (i = 0; i <= db; i++) bb[i] = b[i];
	for (i = db + 1; i < bit[k]; i++) bb[i] = 0;
	
	static int x[524289], y[524289], z[524289];
	NTT(k, aa, x);
	NTT(k, bb, y);
	for (i = 0; i < bit[k]; i++) z[i] = (long long)x[i] * y[i] % Mod;
	NTT_reverse(k, z, cc);
	for (i = 0; i <= da + db; i++) c[i] = (long long)cc[i] * bit_inv[k] % Mod;
}

// Compute the product of two polynomials a[0-da] and b[0-db] naively in O(da * db) time
void prod_poly_naive(int da, int db, int a[], int b[], int c[])
{
	int i, j;
	for (i = 0; i <= da + db; i++) c[i] = 0;
	for (i = 0; i <= da; i++) {
		for (j = 0; j <= db; j++) {
			c[i+j] += (long long)a[i] * b[j] % Mod;
			if (c[i+j] >= Mod) c[i+j] -= Mod;
		}
	}
}

// Compute the product of two polynomials a[0-da] and b[0-db] in an appropriate way
void prod_polynomial(int da, int db, int a[], int b[], int c[])
{
	const int THR = 250000;
	if (THR / (da + 1) >= db + 1) prod_poly_naive(da, db, a, b, c);
	else prod_poly_NTT(da, db, a, b, c);
}

// Compute the remainder obtained by dividing a polynomial a[0-da] by a polynomial b[0-db] using NTT in O(d * log d) time
void mod_poly_NTT(int da, int db, int a[], int b[], int c[])
{
	int i, j, k, l, cur, prev;
	static int x[2][524289], tmp[2][524289];
	if (da < db) {
		for (i = 0; i <= da; i++) c[i] = a[i];
		for (; i < db; i++) c[i] = 0;
		return;
	}

	for (i = 0, j = da; i < j; i++, j--) {
		a[i] ^= a[j];
		a[j] ^= a[i];
		a[i] ^= a[j];
	}
	for (i = 0, j = db; i < j; i++, j--) {
		b[i] ^= b[j];
		b[j] ^= b[i];
		b[i] ^= b[j];
	}
	
	// Compute the inverse of a polynomial b[0-m] mod x^(da - db + 1)
	for (k = 1, cur = 1, prev = 0, x[0][0] = div_mod(1, b[0], Mod); k <= da - db; k <<= 1, cur ^= 1, prev ^= 1) {
		prod_polynomial(k - 1, k - 1, x[prev], x[prev], tmp[0]);
		l = (db < k * 2 - 1)? db: k * 2 - 1;
		for (i = 0; i <= l; i++) tmp[1][i] = b[i];
		prod_polynomial(l, (k - 1) * 2, tmp[1], tmp[0], x[cur]);
		for (i = 0; i < k; i++) {
			x[cur][i] = x[prev][i] * 2 - x[cur][i];
			if (x[cur][i] >= Mod) x[cur][i] -= Mod;
			else if (x[cur][i] < 0) x[cur][i] += Mod;
		}
		for (; i < k * 2; i++) x[cur][i] = Mod - x[cur][i];
	}
	for (i = 0; i <= da - db; i++) tmp[0][i] = a[i];
	prod_polynomial(da - db, da - db, tmp[0], x[prev], x[cur]);

	for (i = 0, j = da; i < j; i++, j--) {
		a[i] ^= a[j];
		a[j] ^= a[i];
		a[i] ^= a[j];
	}
	for (i = 0, j = db; i < j; i++, j--) {
		b[i] ^= b[j];
		b[j] ^= b[i];
		b[i] ^= b[j];
	}
	for (i = 0, j = da - db; i < j; i++, j--) {
		x[cur][i] ^= x[cur][j];
		x[cur][j] ^= x[cur][i];
		x[cur][i] ^= x[cur][j];
	}
	for (i = 0; i <= db; i++) tmp[1][i] = b[i];
	prod_polynomial(db, da - db, tmp[1], x[cur], x[prev]);
	for (i = 0; i <= da; i++) {
		c[i] = a[i] - x[prev][i];
		if (c[i] < 0) c[i] += Mod;
	}
}

// Compute the remainder obtained by dividing a polynomial a[0-da] by a polynomial b[0-db] naively in O(da * db) time
void mod_poly_naive(int da, int db, int a[], int b[], int c[])
{
	int i, j;
	long long tmp;
	for (i = 0; i <= da; i++) c[i] = a[i];
	for (; i < db; i++) c[i] = 0;
	for (i = da - db; i >= 0; i--) {
		tmp = div_mod(c[i+db], b[db], Mod);
		for (j = 0; j <= db; j++) {
			c[i+j] -= tmp * b[j] % Mod;
			if (c[i+j] < 0) c[i+j] += Mod;
		}
	}
}

// Compute the remainder obtained by dividing a polynomial a[0-da] by a polynomial b[0-db] in an appropriate way
void mod_polynomial(int da, int db, int a[], int b[], int c[])
{
	const int THR = 250000;
	if (THR / (da + 1) >= (db + 1)) mod_poly_naive(da, db, a, b, c);
	else mod_poly_NTT(da, db, a, b, c);
}

// Evaluate a polynomial a[0-d] at n points x[1-N] -> y[1-N] in O(d * log d + N * (log N)^2) time
void multipoint_evaluation(int d, int N, int a[], int x[], int y[])
{
	static int i, j, par[524289], deg[524289], *b[524289], *c[524289], head, tail, tmp[3][524289];
	for (i = 1; i <= N; i++) {
		deg[i] = 1;
		b[i] = (int*)malloc(sizeof(int) * (deg[i] + 1));
		c[i] = (int*)malloc(sizeof(int) * deg[i]);
		b[i][0] = (x[i] == 0)? 0: Mod - x[i];
		b[i][1] = 1;
	}
	for (head = 1, tail = N; head < tail; head += 2) {
		par[head] = ++tail;
		par[head+1] = tail;
		deg[tail] = deg[head] + deg[head+1];
		b[tail] = (int*)malloc(sizeof(int) * (deg[tail] + 1));
		c[tail] = (int*)malloc(sizeof(int) * deg[tail]);
		for (i = 0; i <= deg[head]; i++) tmp[0][i] = b[head][i];
		for (i = 0; i <= deg[head+1]; i++) tmp[1][i] = b[head+1][i];
		prod_polynomial(deg[head], deg[head+1], tmp[0], tmp[1], tmp[2]);
		for (i = 0; i <= deg[tail]; i++) b[tail][i] = tmp[2][i];
	}
	mod_polynomial(d, deg[tail], a, b[tail], tmp[0]);
	for (i = 0; i < deg[tail]; i++) c[tail][i] = tmp[0][i];
	for (head = tail - 1; head >= 1; head--) {
		for (d = deg[par[head]] - 1; d > 0 && c[par[head]][d] == 0; d--);
		mod_polynomial(d, deg[head], c[par[head]], b[head], tmp[0]);
		for (i = 0; i < deg[head]; i++) c[head][i] = tmp[0][i];
	}
	for (i = 1; i <= N; i++) y[i] = c[i][0];
	for (head = 1; head <= tail; head++) {
		free(b[head]);
		free(c[head]);
	}
}

long long fact[210000], fact_inv[210000];

long long polynomial_interpolation(int d, int p[], int n)
{
	if (n <= d) return p[n];
	
	int i;
	long long ans = 0, n_fact[2][210000], tmp;
	for (i = 1, n_fact[0][0] = 1, n_fact[1][d] = 1; i <= d; i++) {
		n_fact[0][i] = n_fact[0][i-1] * (n - i + 1) % Mod;
		n_fact[1][d-i] = n_fact[1][d-i+1] * (n - d + i - 1) % Mod;
	}
	for (i = 0; i <= d; i++) {
		if ((i + d) % 2 == 0) ans += p[i] * n_fact[0][i] % Mod * n_fact[1][i] % Mod * fact_inv[d-i] % Mod * fact_inv[i] % Mod;
		else ans += Mod - p[i] * n_fact[0][i] % Mod * n_fact[1][i] % Mod * fact_inv[d-i] % Mod * fact_inv[i] % Mod;
	}
	return ans % Mod;
}

void chmax(int* a, int b)
{
	if (*a < b) *a = b;
}

int main()
{
	int i, N, K, K_max = 0, num[5001] = {};
	long long M;
	scanf("%d %lld", &N, &M);
	M %= Mod;
	for (i = 1; i <= N; i++) {
		scanf("%d", &K);
		num[K]++;
		chmax(&K_max, K);
	}
	
	int j, d = N + K_max, a[210000], b[210000], c[420000], p[210000], mul;
	long long tmp, tmpp, sum;
	for (i = 1, fact[0] = 1; i <= d; i++) fact[i] = fact[i-1] * i % Mod;
	for (i = d - 1, fact_inv[d] = div_mod(1, fact[d], Mod); i >= 0; i--) fact_inv[i] = fact_inv[i+1] * (i + 1) % Mod;
	for (i = 1, tmp = 1, a[0] = 1; i <= d; i++) {
		tmp = tmp * (N - 1 + i) % Mod * fact[i-1] % Mod * fact_inv[i] % Mod;
		a[i] = tmp;
	}
	for (i = 1; i <= d; i++) c[i] = i;
	multipoint_evaluation(K_max, d, num, c, b);
	b[0] = 0;
	/* for (i = 1, b[0] = 0; i <= d; i++) {
		for (j = 1, tmp = i, sum = 0; j <= K_max; j++, tmp = tmp * i % Mod) sum += tmp * num[j];
		b[i] = sum % Mod;
	} */
	prod_polynomial(d, d, a, b, c);
	for (i = 0; i <= d; i++) p[i] = c[i];
	printf("%lld\n", polynomial_interpolation(d, p, M));
	fflush(stdout);
	return 0;
}
0