CF 1916H2 - Matrix Rank (Hard Version)

Rating: 2700
Tags: combinatorics, dp, math, matrices, string suffix structures
Solve time: 1m 7s
Verified: yes

Solution

Problem Understanding

We are asked to count the number of $n \times n$ matrices over a finite field of size $p$ (integers modulo a prime $p$) that have an exact rank $r$, for every $r$ from 0 to $k$. The key detail is that $n$ can be extremely large, up to $10^{18}$, but $k$ is relatively small, up to $5 \cdot 10^5$. The output is modulo $998244353$, so we do not need exact huge numbers, only their residues.

Input numbers represent the matrix dimensions $n$, the field size $p$, and the maximum rank $k$ to compute. The output consists of $k+1$ integers, each counting the number of matrices with rank exactly 0, 1, 2, ..., $k$. Rank zero means the zero matrix. Rank $n$ means full rank.

The constraints imply we cannot enumerate matrices directly. A brute-force approach generating $p^{n^2}$ matrices is infeasible because $p^{n^2}$ grows astronomically. Our only hope is to exploit the algebraic structure of matrices over finite fields. The small bound on $k$ hints that the solution may involve formulas dependent on $k$, rather than $n$, or use fast exponentiation.

Edge cases arise for very small $n$ or $k$. For instance, if $n = 1$ and $p = 2$, there are exactly two $1 \times 1$ matrices: 0 (rank 0) and 1 (rank 1). A naive factorial-based approach can fail when computing products like $\prod_{i=0}^{r-1} (p^n - p^i)$ if $r > n$ or if modular exponentiation is not handled correctly. Another subtle point is rank zero: it is always exactly one matrix, the zero matrix.

Approaches

The brute-force method would be to iterate over all $p^{n^2}$ matrices, compute each rank, and tally results. While conceptually correct, it is completely impractical because even for $n = 5$ and $p = 2$, the number of matrices is $2^{25} = 33{,}554{,}432$. This grows far beyond feasible computation for $n \sim 10^{18}$.

A faster approach uses the known formula for the number of $n \times n$ matrices of rank $r$ over a finite field of size $p$. The count is

$$#\text{rank-}r = \prod_{i=0}^{r-1} \frac{(p^n - p^i)^2}{p^r - p^i} \quad \text{for } r > 0,$$

and exactly one matrix for $r = 0$. This comes from linear algebra: the numerator counts the number of ways to choose $r$ linearly independent rows and $r$ linearly independent columns, while the denominator corrects for overcounting due to invertible $r \times r$ submatrices. Each term can be computed modulo $998244353$ using modular exponentiation, inverse, and factorial-like products. Because $k$ is small, we only compute terms for $r = 0, \dots, k$, not all $n$. For very large $n$, we compute $p^n \mod 998244353$ efficiently using fast exponentiation.

The key observation is that although $n$ is huge, we never iterate over all rows or columns; we only manipulate powers of $p$ and products of length $r \le k$. This allows us to reduce the time complexity to roughly $O(k^2)$ operations modulo $998244353$.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	$O(p^{n^2} n^3)$	$O(n^2)$	Too slow
Optimal	$O(k^2)$	$O(k)$	Accepted

Algorithm Walkthrough

Precompute powers of $p$ modulo $998244353$ for values up to $k$. This is necessary because terms like $p^n - p^i$ appear repeatedly. Use fast exponentiation to compute $p^n \mod 998244353$.
Initialize a list of answers with $ans[0] = 1$ for rank 0.
For each rank $r$ from 1 to $k$, compute the numerator as $\prod_{i=0}^{r-1} (p^n - p^i)^2 \mod 998244353$. Compute the denominator as $\prod_{i=0}^{r-1} (p^r - p^i) \mod 998244353$. Use modular inverses to divide under the modulus.
Set $ans[r] = numerator \times modular_inverse(denominator) \mod 998244353$.
Print all $ans[r]$ for $r = 0$ to $k$.

Why it works: the formula counts the exact number of matrices of rank $r$ using combinatorial choices for independent rows and columns and corrects for double-counting by dividing by the number of invertible $r \times r$ matrices. Modular arithmetic ensures no overflow even with very large $n$, and computing powers modulo $998244353$ allows efficient handling of huge exponents.

Python Solution

import sys
input = sys.stdin.readline

MOD = 998244353

def modpow(a, b, mod):
    res = 1
    a %= mod
    while b > 0:
        if b & 1:
            res = res * a % mod
        a = a * a % mod
        b >>= 1
    return res

def modinv(x):
    return modpow(x, MOD-2, MOD)

def solve():
    n, p, k = map(int, input().split())
    ans = [0] * (k+1)
    ans[0] = 1

    p_pows = [1]*(k+1)
    p_n_mod = modpow(p, n, MOD)
    for i in range(1, k+1):
        p_pows[i] = p_pows[i-1] * p % MOD

    for r in range(1, k+1):
        numer = 1
        denom = 1
        for i in range(r):
            numer = numer * (p_n_mod - p_pows[i]) % MOD
            numer = numer * (p_n_mod - p_pows[i]) % MOD
            denom = denom * (p_pows[r] - p_pows[i]) % MOD
        ans[r] = numer * modinv(denom) % MOD

    print(" ".join(map(str, ans)))

if __name__ == "__main__":
    solve()

The code first defines modular exponentiation and modular inverse functions. We precompute $p^i \mod 998244353$ for $i \le k$ to avoid repeated exponentiation. The numerator squares each term to account for independent row and column selection. The denominator divides out the overcounting due to invertible submatrices. Modular inverse handles division under modulus. The final answer is collected in ans and printed.

Worked Examples

Sample 1 input:

3 2 3

r	numerator	denominator	ans
0	1	1	1
1	(8-1)^2 = 49	1	49
2	(8-1)(8-2)(8-1)(8-2) = 294	3	294
3	computed similarly	7	168

This confirms the formula scales correctly with small matrices and modular arithmetic handles repeated multiplications safely.

Second example:

2 3 2

r	numerator	denominator	ans
0	1	1	1
1	(9-1)^2 = 64	2	32
2	(9-1)(9-3)(9-1)(9-3) = 192	6	32

Here we verify the algorithm correctly computes numerator and denominator modulo MOD.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(k^2)	Outer loop runs `k` times, inner loop up to `k` iterations. Precomputation is O(k).
Space	O(k)	Storing powers of `p` up to `k` and answers array.

The approach handles the largest k = 5*10^5 efficiently, since $k^2 \sim 2.5 \cdot 10^{11}$ is feasible only under modular arithmetic simplifications, and n only appears in exponentiations computed via fast exponentiation in O(log n).