Nauuo and Pictures (hard version)
Contest: Codeforces Round 564 (Div. 1)
Difficulty rating: 2600
Tags: dp, probabilities

The only difference between easy and hard versions is constraints.

Nauuo is a girl who loves random picture websites.

One day she made a random picture website by herself which includes $$$n$$$ pictures.

When Nauuo visits the website, she sees exactly one picture. The website does not display each picture with equal probability. The $$$i$$$-th picture has a non-negative weight $$$w_i$$$, and the probability of the $$$i$$$-th picture being displayed is $$$\frac{w_i}{\sum_{j=1}^nw_j}$$$. That is to say, the probability of a picture to be displayed is proportional to its weight.

However, Nauuo discovered that some pictures she does not like were displayed too often.

To solve this problem, she came up with a great idea: when she saw a picture she likes, she would add $$$1$$$ to its weight; otherwise, she would subtract $$$1$$$ from its weight.

Nauuo will visit the website $$$m$$$ times. She wants to know the expected weight of each picture after all the $$$m$$$ visits modulo $$$998244353$$$. Can you help her?

The expected weight of the $$$i$$$-th picture can be denoted by $$$\frac {q_i} {p_i}$$$ where $$$\gcd(p_i,q_i)=1$$$, you need to print an integer $$$r_i$$$ satisfying $$$0\le r_i<998244353$$$ and $$$r_i\cdot p_i\equiv q_i\pmod{998244353}$$$. It can be proved that such $$$r_i$$$ exists and is unique.

Input format:
The first line contains two integers $$$n$$$ and $$$m$$$ ($$$1\le n\le 2\cdot 10^5$$$, $$$1\le m\le 3000$$$) — the number of pictures and the number of visits to the website.

The second line contains $$$n$$$ integers $$$a_1,a_2,\ldots,a_n$$$ ($$$a_i$$$ is either $$$0$$$ or $$$1$$$) — if $$$a_i=0$$$ , Nauuo does not like the $$$i$$$-th picture; otherwise Nauuo likes the $$$i$$$-th picture. It is guaranteed that there is at least one picture which Nauuo likes.

The third line contains $$$n$$$ positive integers $$$w_1,w_2,\ldots,w_n$$$ ($$$w_i \geq 1$$$) — the initial weights of the pictures. It is guaranteed that the sum of all the initial weights does not exceed $$$998244352-m$$$.

Output format:
The output contains $$$n$$$ integers $$$r_1,r_2,\ldots,r_n$$$ — the expected weights modulo $$$998244353$$$.

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

Note:
In the first example, if the only visit shows the first picture with a probability of $$$\frac 2 3$$$, the final weights are $$$(1,1)$$$; if the only visit shows the second picture with a probability of $$$\frac1 3$$$, the final weights are $$$(2,2)$$$.

So, both expected weights are $$$\frac2 3\cdot 1+\frac 1 3\cdot 2=\frac4 3$$$ .

Because $$$332748119\cdot 3\equiv 4\pmod{998244353}$$$, you need to print $$$332748119$$$ instead of $$$\frac4 3$$$ or $$$1.3333333333$$$.

In the second example, there is only one picture which Nauuo likes, so every time Nauuo visits the website, $$$w_1$$$ will be increased by $$$1$$$.

So, the expected weight is $$$1+2=3$$$.

Nauuo is very naughty so she didn't give you any hint of the third example.

Example input:
2 1
0 1
2 1

Example output:
332748119
332748119

Example input:
1 2
1
1

Example output:
3

Editorial:
First, let's focus on a single picture with weight $$$w$$$ which Nauuo likes, so we only have to know the sum of the weights of the pictures Nauuo likes ($$$SA=\sum\limits_{i=1}^nw_i[a_i=1]$$$) and the sum of the disliked ones ($$$SB=\sum\limits_{i=1}^nw_i[a_i=0]$$$) instead of all the $$$n$$$ weights.
Then, we can use DP to solve this problem.
Let $$$f_w[i][j][k]$$$ be the expected weight of a picture Nauuo likes with weight $$$w$$$ after $$$i$$$ visits when $$$SA=j$$$ and $$$SB=k$$$ now.
Obviously, $$$f_w[0][j][k]=w$$$.
The state transition:
• The next visit displays the picture we focus on. Probaility: $$$\frac w{j+k}$$$. Lead to: $$$f_{w+1}[i-1][j+1][k]$$$.
• The next visit displays a picture Nauuo likes but is not the one we focus on. Probaility: $$$\frac{j-w}{j+k}$$$. Lead to: $$$f_w[i-1][j+1][k]$$$.
• The next visit displays a picture Nauuo doesn't like. Probaility: $$$\frac k{j+k}$$$. Lead to: $$$f_w[i-1][j][k-1]$$$.
So, $$$f_w[i][j][k]=\frac w{j+k}f_{w+1}[i-1][j+1][k]+\frac{j-w}{j+k}f_w[i-1][j+1][k]+\frac k{j+k}f_w[i-1][j][k-1]$$$.
Let $$$g_w[i][j][k]$$$ be the expected weight of a picture Nauuo doesn't like with weight $$$w$$$ after $$$i$$$ visits when $$$SA=j$$$ and $$$SB=k$$$ now. The state transition is similar.
Note that $$$i,\,j,\,k,\,m$$$ have some relation. In fact we can let $$$f'_w[i][j]$$$ be $$$f_w[m-i-j][SA+i][SB-j]$$$ ($$$SA$$$ and $$$SB$$$ are the initial ones here).
But up to now, we can only solve the easy version.
To solve the hard version, let's introduce a lemma:
$$$$$$f_w[i][j][k]=wf_1[i][j][k]$$$$$$
Proof:
Obviously, this is true when $$$i=0$$$.
Then, suppose we have already proved $$$f_w[i-1][j][k]=wf_1[i-1][j][k]$$$.
$$$\begin{aligned}f_1[i][j][k]&=\frac 1{j+k}f_2[i-1][j+1][k]+\frac{j-1}{j+k}f_1[i-1][j+1][k]+\frac k{j+k}f_1[i-1][j][k-1]\\&=\frac2{j+k}f_1[i-1][j+1][k]+\frac{j-1}{j+k}f_1[i-1][j+1][k]+\frac k{j+k}f_1[i-1][j][k-1]\\&=\frac{j+1}{j+k}f_1[i-1][j+1][k]+\frac k{j+k}f_1[i-1][j][k-1]\end{aligned}$$$
$$$\begin{aligned}f_w[i][j][k]&=\frac w{j+k}f_{w+1}[i-1][j+1][k]+\frac{j-w}{j+k}f_w[i-1][j+1][k]+\frac k{j+k}f_w[i-1][j][k-1]\\&=\frac{w(w+1)}{j+k}f_1[i-1][j+1][k]+\frac{w(j-w)}{j+k}f_1[i-1][j+1][k]+\frac {wk}{j+k}f_1[i-1][j][k-1]\\&=\frac{w(j+1)}{j+k}f_1[i-1][j+1][k]+\frac {wk}{j+k}f_1[i-1][j][k-1]\\&=wf_1[i][j][k]\end{aligned}$$$
Also, a brief but not so strict proof: the increment in each step is proportional to the expectation.
So, we only have to calculate $$$f_1[i][j][k]$$$ ($$$f'_1[i][j]$$$).
In conclusion:
$$$f'_1[i][j]=1\ (i+j=m)$$$
$$$f'_1[i][j]=\frac{SA+i+1}{SA+SB+i-j}f'_1[i+1][j]+\frac{SB-j}{SA+SB+i-j}f'_1[i][j+1]\ (i+j<m)$$$
$$$g'_1[i][j]=1\ (i+j=m)$$$
$$$g'_1[i][j]=\frac{SA+i}{SA+SB+i-j}g'_1[i+1][j]+\frac{SB-j-1}{SA+SB+i-j}g'_1[i][j+1]\ (i+j<m)$$$
If $$$a_i=1$$$, the expected weight of the $$$i$$$-th picture is $$$w_if'_1[0][0]$$$, otherwise, the expected weight is $$$w_ig'_1[0][0]$$$.
Last question: how to calculate the result modulo $$$998244353$$$?
If you don't know how, please read the wiki to learn it.
You can calculate and store all the $$$\mathcal O(m)$$$ inverses at first, then you can get an $$$\mathcal O(n+m^2+m\log p)$$$ solution instead of $$$\mathcal O(n+m^2\log p)$$$ ($$$p=998244353$$$ here).