卢卡斯定理优化

P5598【XR-4】混乱度

题目描述

小 X 有 $n$ 种颜色的球，其中第 $i$ 种颜色的球共有 $a_i$ 个，同色的球无法区分。定义第 $l \sim r$ 种颜色的混乱度 $f(l, r)$ 为：将第 $l \sim r$ 种颜色的所有球排成一排，总共的方案数对 $p$ 取模后的值。小 X 想请你帮忙计算下列式子的值：

$\sum_{l=1}^n \sum_{r=l}^n f(l, r)$

题目分析

一眼就看出来f是一个可重集合的排列，而且计算很快，困难的是a的值域太大了，一般的卢卡斯定理还不能通过此题。所以要进行预处理来加快卢卡斯定理的计算速度。

代码

//
// Created by mrx on 2022/10/6.
//

#include <vector>
#include <algorithm>
#include <iostream>
#include <array>

using ll = long long;
int P = 1;

struct digit {
	std::vector<int> d;
	int cnt;

	digit() : d(10), cnt(0) {}

	digit(ll x) : d(10), cnt(0) {
		ll tmp = x;
		while (tmp) {
			d[cnt++] = tmp % P;
			tmp /= P;
		}
	}

	void add(const digit& rhs) {
		for (int i = 0; i < std::max(cnt, rhs.cnt); i++) {
			d[i] += rhs.d[i];
			if (d[i] >= P) {
				d[i] -= P;
				d[i + 1]++;
			}
		}
	}
};


int main() {
	std::ios::sync_with_stdio(false);
	std::cin.tie(nullptr);
	std::cout.tie(nullptr);
	int n, p;
	std::cin >> n >> p;
	P = 1;
	while (P < 5000)P *= p;
	P /= p;
	std::vector<std::vector<int>> C(P + 1, std::vector<int>(P + 1));
	C[0][0] = 1;
	for (int i = 1; i <= P; ++i) {
		C[i][0] = C[i][i] = 1;
		for (int j = 1; j < P; ++j) {
			C[i][j] = (C[i - 1][j] + C[i - 1][j - 1]) % p;
		}
	}
	auto lucas = [&](const digit& n, const digit& m) {
		int s = 1;
		for (int i = 0; i < m.cnt; ++i) {
			s *= C[n.d[i]][m.d[i]] %= p;
		}
		return s;
	};
	std::vector<digit> a(n + 1);
	for (int i = 0; i < n; ++i) {
		ll x;
		std::cin >> x;
		a[i] = digit(x);
	}

	std::vector<int> r(n + 1);
	a[n].cnt = 1;
	int st = 0;
	for (int i = 0; i <= n; i++) {
		if (a[i].cnt) {
			for (int j = i - 1; j >= st; j--)r[j] = i;
			st = i;
		}
	}

	ll ans = 0;
	for (int i = 0; i < n; ++i) {
		int ptr = i;
		digit tmp = a[ptr];
		int cur = 1;
		while (cur && ptr < n) {
			ans += (r[ptr] - ptr) * cur;
			ptr = r[ptr];
			tmp.add(a[ptr]);
			cur *= lucas(tmp, a[ptr]);
			cur %= p;
		}
	}
	std::cout << ans << '\n';
	return 0;
}