Matrix Rank (Hard Version)
Contest: Good Bye 2023
Difficulty rating: 2700
Tags: combinatorics, dp, math, matrices, string suffix structures

This is the hard version of the problem. The only differences between the two versions of this problem are the constraints on $$$k$$$. You can make hacks only if all versions of the problem are solved.

You are given integers $$$n$$$, $$$p$$$ and $$$k$$$. $$$p$$$ is guaranteed to be a prime number.

For each $$$r$$$ from $$$0$$$ to $$$k$$$, find the number of $$$n \times n$$$ matrices $$$A$$$ of the field$$$^\dagger$$$ of integers modulo $$$p$$$ such that the rank$$$^\ddagger$$$ of $$$A$$$ is exactly $$$r$$$. Since these values are big, you are only required to output them modulo $$$998\,244\,353$$$.

$$$^\dagger$$$ https://en.wikipedia.org/wiki/Field_(mathematics)

$$$^\ddagger$$$ https://en.wikipedia.org/wiki/Rank_(linear_algebra)

Input format:
The first line of input contains three integers $$$n$$$, $$$p$$$ and $$$k$$$ ($$$1 \leq n \leq 10^{18}$$$, $$$2 \leq p < 998\,244\,353$$$, $$$0 \leq k \leq 5 \cdot 10^5$$$).

It is guaranteed that $$$p$$$ is a prime number.

Output format:
Output $$$k+1$$$ integers, the answers for each $$$r$$$ from $$$0$$$ to $$$k$$$.

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

Example input:
3 2 3

Example output:
1 49 294 168

Example input:
1 51549919 2

Example output:
1 51549918 0

Editorial:
Let $$$dp_{m,k}$$$ denote the number of $$$m \times n$$$ matrices whose rank is $$$k$$$.
The recurrence for $$$dp_{m,k} = dp_{m-1,k} \cdot p^k + dp_{m-1,k-1} \cdot (p^n-p^{k-1})$$$ with the base case of $$$dp_{0,0}=1$$$.
Consider enumerating $$$m \times n$$$ matrices $$$A$$$ such that $$$\text{rank}(A)=k$$$ based on a $$$(m-1) \times n$$$ matrix $$$B$$$ and a row vector $$$C$$$ such that $$$A$$$ is made by appending $$$C$$$ to $$$B$$$:
• If $$$\text{rank}(B) = \text{rank}(A)$$$, then $$$C$$$ must be inside the span of $$$B$$$ which has size $$$p^k$$$.
• If $$$\text{rank}(B)+1=\text{rank}(A)$$$, then $$$C$$$ must be outside the span of $$$B$$$ which has size $$$p^n-p^{k-1}$$$.
The $$$p^n-p^{k-1}$$$ term in our recurrence is annoying to handle. So we will define $$$g_{m,k}$$$ as $$$g_{m,k} = g_{m-1,k} \cdot p^k + g_{m-1,k-1}$$$. It should be easy to see that $$$dp_{m,k}=g_{m,k} \cdot \prod\limits_{i=0}^{k-1} (p^n-p^i)$$$.
Method 1: Easy Version Only
We will show how to compute $$$g_{2m}$$$ from $$$g_m$$$. I claim that $$$g_{2m,k} = \sum\limits_{a+b=k} g_{m,a} \cdot g_{m,b} \cdot p^{a(m-b)}$$$.
The justification for this is to consider moving the DP state from $$$(0,0) \to (m,a) \to (2m,a+b)$$$.
The contribution of moving $$$(m,a) \to (2m,a+b)$$$ is actually very similar to that of moving $$$(0,a) \to (m,a+b)$$$, except every time we multiply the value by $$$p^k$$$, we have to instead multiply it by $$$p^{k+a}$$$. And conveniently, the number of times we need to additionally multiply by $$$p^k$$$ is $$$m-b$$$.
To convolute these sequences, we will use a technique that is similar to chirp-z transform. Note that $$$ab = \binom{a+b}{2} - \binom{a}{2} - \binom{b}{2}$$$, you can prove this geometrically by looking at $$$\binom{x}{2}$$$ as a triangle pyramid.
$$$$$$\begin{aligned} \sum\limits_{a+b=k} g_{m,a} \cdot g_{m,b} \cdot p^{a(m-b)} &= \sum\limits_{a+b=k} g_{m,a} \cdot g_{m,b} \cdot p^{am + \binom{a}{2} + \binom{b}{2} - \binom{a+b}{2}} \\ &= p^{-\binom{k}{2}} \sum\limits_{a+b=k} (g_{m,a} \cdot p^{am + \binom{a}{2}}) \cdot (g_{m,b} \cdot p^{\binom{b}{2}}) \end{aligned}$$$$$$
Therefore, it is we can get $$$g_{2m}$$$ from $$$g_m$$$ by convoluting $$$(g_{m,i} \cdot p^{im+\binom{i}{2}})$$$ and $$$(g_{m,i} \cdot p^{\binom{i}{2}})$$$ in $$$O(k \log k)$$$ time. By performing a process similar to fast exponentiation, we can compute $$$g_n$$$ in $$$O(k \log k \log n)$$$.
Method 2: Easy Version Only
Consider the ordinary generating function $$$G_k = \sum\limits g_{m,k} x^m$$$. Based on the recurrence above, We have $$$G_k = x(p^kG_k + G_{k-1})$$$. This gives us $$$G_k = \frac{xG_{k-1}}{1-p^kx} = x^k \prod\limits_{i=0}^k \frac{1}{1-p^ix}$$$.
We want to find $$$[x^n] G_0, [x^n] G_1, \ldots, [x^n] G_k$$$. Notice that $$$[x^n] G_k = [x^{n-k}] \prod\limits_{i=0}^k \frac{1}{1-p^ix} = [x^k] x^n ~\%~ (\prod\limits_{i=0}^k x-p^i)$$$. Where $$$A ~\%~ B$$$ denotes the unique polynomial $$$C$$$ with $$$\text{deg}(C)<\text{deg}(B)$$$ and there exists $$$D$$$ such that $$$A = BD + C$$$.
We can compute all required coefficient in a divide and conquer-esque way. Note that $$$A ~\%~ B = (A ~\%~ BC)~\%~B$$$.
We want $$$\text{dnc}(l,r)$$$ to find all the answers for $$$[x^n] G_l, [x^n] G_{l+1}, \ldots, [x^n] G_r$$$. $$$\text{dnc}(l,r)$$$ will be provided with $$$P(x)=x^n ~\%~ (\prod\limits_{i=0}^r x-p^i)$$$ and $$$Q(x) = \prod\limits_{i=0}^{l} x-p^i$$$.
Then the pseudocode of the DnC is as follows:
The above code is at least quadratic. However, we only need to save the last $$$(r-l+1)$$$ non-zero coefficients from both $$$P$$$ and $$$Q$$$. If implemented properly, this code runs in $$$O(k \log k \log n)$$$.
Method 3: Hard Version
As with the first part of method $$$2$$$, we want to find $$$[x^n] G_0, [x^n] G_1, \ldots, [x^n] G_k$$$ where $$$G_k = x^k \prod\limits_{i=0}^k \frac{1}{1-p^ix}$$$. We will try to decompose $$$\prod\limits_{i=0}^k \frac{1}{1-p^ix}$$$ into partial fractions.
Repeated factors in partial fractions is hard to deal with. It turns out, we do not need to deal with them at all. Suppose $$$ord$$$ is the smallest positive integer that such $$$p^{ord} = 1 \pmod{998244353}$$$. Then $$$p^0,p^1,\ldots,p^{ord-1}$$$ are all distinct. Furthermore, $$$dp_{n,k}=0$$$ for $$$k \geq ord$$$ as $$$dp_{m,k}=g_{m,k} \cdot \prod\limits_{i=0}^{k-1} (p^n-p^i)$$$.
Therefore, we can safely assume $$$k < ord$$$ so that $$$p^0,p^1,\ldots,p^k$$$ will be all distinct.
By using the cover-up rule on $$$\prod\limits_{i=0}^k \frac{1}{1-p^ix}$$$, we obtain a partial fraction decomposition of $$$\sum\limits_{i=0}^k \frac{c_i}{1-p^ix}$$$ where $$$c_i = \prod\limits_{j=0, i\neq j}^k \frac{1}{1-p^{j-i}}$$$.
Let $$$L_i=\prod_{j=1}^i\frac 1{1-p^{-j}},R_i=\prod_{j=1}^i\frac 1{1-p^j}$$$，so that $$$c_i=L_iR_{k-i}$$$.
$$$$$$\begin{aligned}[x^n] G_k &= [x^{n-k}] \prod\limits_{i=0}^k \frac{1}{1-p^ix} \\ &= [x^{n-k}] \sum\limits_{a+b=k} \frac{L_a R_b}{1-p^ax} \\ &= \sum\limits_{a+b=k} p^{a(n-k)} L_a R_b \\ &=\sum\limits_{a+b=k} p^{\binom{n-b}{2} - \binom{n-k}{2} - \binom{a}{2}} L_a R_b \\ &=p^{-\binom{n-k}{2}} \sum\limits_{a+b=k} (L_a \cdot p^{-\binom{a}{2}}) (R_{b} \cdot p^{\binom{n-b}{2}}) \end{aligned}$$$$$$
So we can find $$$[x^n] G_0, [x^n] G_1, \ldots, [x^n] G_k$$$ by convoluting $$$(L_i \cdot p^{\binom{i}{2}})$$$ and $$$(R_{i} \cdot p^{\binom{n-i}{2}})$$$.
The final time complexity is $$$O(k \log k)$$$.
Method 4: Hard Version
Thanks to Endagorion for finding this solution during testing.
For an $$$n$$$-dimensional vector space of $$$\mathbb{F}_q$$$, the number of $$$k$$$-dimensional vector subspace is denoted as $$$\binom{n}{k}_q$$$ and is known as the q-analogue of binomial coefficients.
It is known that $$$\binom{n}{k}_q = \frac{(q^n-1)\ldots(q^{n-k+1}-1)}{(1-q)\ldots(1-q^k)}$$$. This fact is not too hard to derive by considering the number of possible independent sets dividing the the number of sets that results in the same span. Here it is helpful to use the form $$$\binom{n}{k}_q = \frac{[n!]_q}{[k!]_q[n-k!]_q}$$$ where $$$[n!]_q = (1) \cdot (1+q) \ldots \cdot (1+q+\ldots+q^{n-1})$$$. If you want to read more about these series, I recommend checking out Enumerative Combinatorics.
We will proceed with PIE (principle of inclusion-exclusion) on these $$$q$$$-binomials. We want to find the number of ways to choose $$$n$$$ vectors that spans exactly some $$$k$$$-dimensional vector space, but this is hard. However, we realize that we are able to count the number of ways to choose to choose $$$n$$$ vectors that spans a subset of some $$$k$$$-dimensional vector space. Let $$$f_k$$$ and $$$g_k$$$ denote these values respectively.
Then we have $$$g_k = \sum\limits_{i=0}^k f_i \cdot \binom{k}{i}_q$$$. Note that this is very similar to egf-convolution with the identity sequence. Indeed, if we define the generating functions $$$F = \sum f_k \frac{x^k}{[k!]_q}$$$, $$$G = \sum g_k \frac{x^k}{[k!]_q}$$$ and the identity sequence $$$I = \sum \frac{x^k}{[k!]_q}$$$, we have $$$F \cdot I = G$$$.
We are able to calculate $$$I$$$ and $$$G$$$, so we can find $$$F = G \cdot I^{-1}$$$.
Then, the number of matrices with rank $$$k$$$ is simply $$$f_k \cdot \binom{n}{k}_q$$$.
Of course, a caveat is that division might not be defined when some $$$[n!]_q=0$$$, and it is a simple exercise to the reader to figure out why defining $$$0^{-1}=0$$$, which aligns with our usual definition of $$$a^{MOD-2}=a^{-1}$$$ does not break anything in the solution.
The final time complexity is $$$O(k \log k)$$$.