Plurm's blog

By Plurm, history, 7 months ago, In English

Disclaimer: This might enter the realm of mathematical trickery, and not competitive programming itself, but I think it's beautiful enough to mention it here in Codeforces, and maybe invite some competitive programmers to somewhere else where I've recently moved to...

Problem

Given an integer $$$0 \leq N \leq 10^9$$$ as input. Count how many $$$8$$$-tuple $$$x \in \mathbb{Z}^8$$$ are there such that $$$x_1^2 + \dots + x_8^2 = N$$$. The answer may be large, so just output the answer modulo $$$10^9+7$$$.

Ideas

Obvious brute-force takes $$$O(N^8)$$$. Observing that iterating in the square root range is enough reduces it to $$$O(\sqrt{N}^8) = O(N^4)$$$. Trying to divide the problems into two $$$4$$$-tuples with $$$x_1^2+x_2^2+x_3^2+x_4^2 = M$$$ and $$$x_5^2+x_6^2+x_7^2+x_8^2 = N-M$$$ reduces to $$$O(N \cdot \sqrt{N}^4) = O(N^3)$$$. Reformulating the problem into DP/Knapsack gives an $$$O(N^{\frac{3}{2}})$$$ algorithm. Can we do better?

It turns out that in this case, there is an explicit formula! (Yes, mathematical trickery, I know) But the beauty is not in the formula itself, but rather "how does one discover this kind of formula".

Spoiler

I've written a full article on this, which, I feel like it's too long to be rewritten as a Codeforces blog, so I'm putting the link here: My Article.

Note that I've also submitted this article to SoME4, summer 2025, even though I don't think it's a good exposition--it's more like a self-study notes to actually go through and internalize what I've learned.

Please let me know if there is any error or mistake. All feedbacks are appreciated (not just limited to error reporting). I'd be happy to see discussions around this topic!

Also, if I have time, I would maybe set up an actual (experimental -- only for fun) competitive programming task generalizing this (I'm not sure if anyone has prepared a task like this before -- I'm not going to take credits, as it is very classical in the theory and literature, just that it's maybe a bit far from our usual competitive programming style):

(Generalized) Problem

Given integers $$$2 \leq L \leq 1\,000$$$, $$$0 \leq N \leq 10^9$$$, and $$$1 \leq k \leq 100$$$ as input. We guarantee that $$$N$$$ has no prime divisor greater than $$$L$$$. Count the number of integer $$$k$$$-tuples $$$(x_1, \dots, x_k) \in \mathbb{Z}^k$$$ such that $$$\sum_{i=1}^k x_i^2 = N$$$. Since the answer can be large, output modulo $$$10^9+7$$$.

Spoiler
  • Vote: I like it
  • +51
  • Vote: I do not like it

»
7 months ago, hide # |
 
Vote: I like it +20 Vote: I do not like it

For the original problem consider the polynomial $$$P(x) = \sum\limits_{i=1}^{\sqrt(N)} x^{i^2}$$$. Then you just need to evaluate coefficient of $$$x^N$$$ in $$$P(x)^8$$$ which can be done with FFT. However, the problem is still not solved yet. Since $$$x \in Z^8$$$, you can have $$$2^8$$$ possibilities for each solutions. Also you can have a zero also, so I think the solution will be coefficient of $$$x^N$$$ in $$$\sum\limits_{j=0}^{8} \binom{8}{8 - j} \cdot P(x)^j \cdot 2^j$$$. basically we take $$$(8 - j)$$$ zeroes and then remaining non-zero solutions and $$$2^j$$$ for the sign. This sum is just $$$(1+2\cdot P(x))^8$$$ but this is feasible till $$$N \leq 10^6$$$.

The complexity is $$$O(N \cdot log(N))$$$, which is better than $$$O(N^{\frac{3}{2}})$$$.

  • »
    »
    7 months ago, hide # ^ |
     
    Vote: I like it 0 Vote: I do not like it

    Excellent! This looks like another elegant approach using FFT. Yes, it would be better than the classic DP approach. Indeed, the coefficients of $$$(1+2P(x))^8$$$ are exactly the answers and it takes time $$$O(N \log N)$$$, better than $$$O(N^{\frac{3}{2}})$$$.

    In general I also think it would generalize pretty well: $$$(1+2P(x))^k$$$ where $$$P(x) := \sum_{i=1}^{\sqrt{N}} x^{i^2}$$$ should work even for the generalized problem, where I think (if I'm not mistaken) it probably takes $$$O(N \log N + N \log k)$$$ (do FFT for $$$P$$$ and exponentiate the result by $$$k$$$, then inverse FFT back to read off the coefficients).

»
7 months ago, hide # |
 
Vote: I like it 0 Vote: I do not like it

Thanks!

It was a nice read for me