Products in ML-KEM
In ML-KEM polynomials live in the ring Z3329[x] / (x²⁵⁶ + 1): the coefficients are remainders on division by 3329, and x²⁵⁶ is replaced by −1. Given a polynomial a and many second factors b1, …, bq. Find each product c = a · bi in this ring.
To keep the input short, bi is given by two numbers s and t: its coefficient of xj is (s · j² + t · j + s + t) mod 3329, j = 0 … 255. And to keep the output short, instead of the 256 coefficients c0, …, c255 (each from 0 to 3328) print Σ (j + 1) · cj mod (10⁹ + 7).
Input
The first line holds the 256 coefficients of a, from the constant term, each from 0 to 3328. The second holds q (1 ≤ q ≤ 5000). Then q lines: s and t (0 ≤ s, t < 3329).
Output
For each factor, the checksum of the product on a line of its own.
Examples
Example 1
0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 2 1 0 0 1
53631721 5595393
Example 2
1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 2 3
56233561
Limits
- Time · Python
- 18500 ms
- Time · Go
- 1000 ms
- Memory
- 256 MB
- Judging
- 7 tests
Walkthrough
Open it step by step: each step gives away a little more. You can jump straight to the solution, but a problem solved with one hint stays with you longer.
1Hint 1
By the schoolbook method one product is 65,536 multiplications, and there are five thousand of them. The polynomial a is the same in all products: what work on it can be done once?
2Hint 2
17 modulo 3329 is a root of unity of order 256, and 17¹²⁸ ≡ −1. So x²⁵⁶ + 1 = x²⁵⁶ − 17¹²⁸ factors, and multiplication modulo each factor is cheap.
3The idea
The incomplete number-theoretic transform of ML-KEM: seven layers of butterflies turn a polynomial into 128 remainders modulo x² − γi, where products are computed on pairs of coefficients, and the inverse transform assembles the answer. The transform of a is done once; for each bi, a forward transform, 128 small products and an inverse.
4Walkthrough
The tables ζk = 17BitRev7(k) and γk = 172·BitRev7(k)+1 modulo 3329 are built once. The forward transform is lengths 128, 64, …, 2, with the butterfly t = ζ · f[j + len], f[j] ± t; the inverse is the same butterflies in reverse order and multiplication by 3303 = 128⁻¹. In the middle (a0 + a1x)(b0 + b1x) = a0b0 + a1b1γ + (a0b1 + a1b0)x.
4+Solution
5Complexity
Per product: 896 multiplications in the forward transform, 640 in the middle and 1152 in the inverse, instead of 65,536.
6Common mistakes
- The schoolbook method. This is
wrong.py: it makes it on the examples, but five thousand factors are more than three hundred million multiplications. - An ordinary cyclic convolution. If you forget that x²⁵⁶ = −1 and add what spills over the edge with a plus, you get a product in a different ring. The seventh test multiplies by x²⁵⁵, and the sign decides everything there.
- A complete transform. There is no eighth layer: it would need a root of order 512, and 3328 = 13 · 2⁸ is not divisible by 512.
- Negative remainders. In Go and C the difference (f[j] − t) % 3329 can be negative: add 3329.