Count Voting
Contest: Codeforces Round 854 by cybercats (Div. 1 + Div. 2)
Difficulty rating: 2600
Tags: combinatorics, dp, math

There are $$$n$$$ people that will participate in voting. Each person has exactly one vote.

$$$i$$$-th person has a team $$$t_i$$$ ($$$1 \leq t_i \leq n$$$) where $$$t_i = t_j$$$ means $$$i$$$, $$$j$$$ are in the same team. By the rules each person should vote for the person from the different team. Note that it automatically means that each person can't vote for himself.

Each person knows the number of votes $$$c_i$$$ he wants to get. How many possible votings exists, such that each person will get the desired number of votes? Due to this number can be big, find it by modulo $$$998\,244\,353$$$.

Input format:
The first line contains a single integer $$$n$$$ ($$$1 \leq n \leq 200$$$) — the number of people.

The second line contains $$$n$$$ integers $$$c_1, c_2, \ldots, c_n$$$ ($$$0 \leq c_i \leq n$$$) — desired number of votes. It is guaranteed, that $$$\sum\limits_{i=1}^{n} c_i = n$$$.

The third line contains $$$n$$$ integers $$$t_1, t_2, \ldots, t_n$$$ ($$$1 \leq t_i \leq n$$$) — team numbers.

Output format:
Print a single integer — the number of possible votings by modulo $$$998\,244\,353$$$.

Constraints:
Time limit: 1.0 s
Memory limit: 256.0 MB

Note:
In the first test there are two possible votings: $$$(2, 3, 1)$$$, $$$(3, 1, 2)$$$.

In the third test there are five possible votings: $$$(3, 3, 2, 2, 1)$$$, $$$(2, 3, 2, 3, 1)$$$, $$$(3, 3, 1, 2, 2)$$$, $$$(3, 1, 2, 3, 2)$$$, $$$(2, 3, 1, 3, 2)$$$.

Example input:
3
1 1 1
1 2 3

Example output:
2

Example input:
5
2 0 1 0 2
1 2 3 4 5

Example output:
10

Editorial:
Let's solve the problem using inclusion-conclusion principle. If we do not consider teams, there are $$$\frac{n!}{\prod\limits_{i=1}^{n} c_i!}$$$ ways to make votings.
But now consider such sets of bad votings: person $$$i$$$ made a vote for the person from the same team $$$t_i$$$. We want to calculate the size of the union of such sets of votings.
So the answer to the problem is $$$\sum\limits_{k=0}^{n} (-1)^k f_k$$$, where $$$f_k$$$ is the number of ways to make votings, where we firstly fix some subset of $$$k$$$ people who will vote for the same team (for others the vote can be any) and calculate the number of votings for it, and we sum these numbers for all subsets of $$$k$$$ people.
Let's define $$$b_1, b_2, \ldots, b_n$$$, where $$$0 \leq b_i \leq c_i$$$ is the number of people from the same team who made a vote for $$$i$$$-th person (so, $$$\sum\limits_{i=1}^{n} b_i = k$$$). So, $$$f_k = \sum\limits_{0 \leq b_i \leq c_i} cnt(b_1, b_2, \ldots, b_n)$$$, where $$$cnt(b_1, b_2, \ldots, b_n)$$$ is the number of votings, where for each $$$i$$$ we firstly fix $$$b_i$$$ people from team $$$t_i$$$ that will vote for $$$i$$$ and calculate the number of votings, and we sum these numbers for all ways to fix.
How $$$cnt$$$ is calculated? There are $$$\frac{(\sum\limits_{i=1}^{n} c_i - b_i)!}{\prod\limits_{i=1}^{n} (c_i - b_i)!} = \frac{(n - k)!}{\prod\limits_{i=1}^{n} (c_i - b_i)!}$$$ ways to make votes for people who were not selected into subset. This should be multiplied for the number of ways to vote for people that were fixed. This should be found for each team separately.
Let's consider some team $$$1, 2, \ldots, m$$$ of $$$m$$$ people (WLOG they are first $$$m$$$ people). There are $$$\frac{m!}{(\prod\limits_{i=1}^{m} b_i!) (m - \sum\limits_{i=1}^{m} b_i)!}$$$ ways (*) to fix people in this team. These counts should be multiplied for all teams.
Now to sum all $$$cnt$$$ for all $$$b$$$ we of course can't iterate over all possible $$$b$$$, we will do it using dynamic programming. Let's iterate over all possible teams and calculate $$$dp_{i,k}$$$ — the sum of current counts for all prefixes of arrays $$$b$$$, where we considered first $$$i$$$ teams and the sum of elements of $$$b$$$ on this prefix is $$$k$$$.
When we add a new team to our $$$dp$$$, we can write $$$dp_{i+1, k+s} += dp_{i,k} t_{i+1,s}$$$, where $$$t_{i+1,s}$$$ is the sum of (*) multiplied by $$$\frac{1}{\prod\limits_{i=1}^{m} (c_i - b_i)!}$$$ for all ways to choose $$$b$$$ inside the team $$$i+1$$$ with sum of $$$b_j$$$ in this team equal to $$$s$$$.
To calculate values $$$t_{i,s}$$$ for each team $$$i$$$ we can also use a prefix dynamic programming inside the team.
At the end we found $$$f_k = dp_{t,k} (n-k)!$$$, where $$$t$$$ is the number of teams.
So, the time complexity of this solution is $$$O(n^2)$$$. Where are many different ways of how the solution can be implemented (and even optimized with power series), but the inclusion-conclusion is necessary to solve the problem.