Finding the sum of four squares

There is Lagrange's four-square theorem, which statea that


Given an integer $N$, we can write $N$ as a sum of four squares
$A^2+B^2+C^2+D^2=N$.

How can I find a valid solution with time complexity of $O(1)$ or $O(\sqrt N)$? I only have to find one solution.

Here is my $O(1)$ algorithm:

while(count!=4){
     solution = sqrt(N)
     print solution
     N-= solution*solution
     count++       
}

But it does not works for all $N$, for example it fails for $N=23$. Is there any other solution?

Try this: Legendre's_three-square_theorem

Assume $4\not\mid N$. If not, divide $N$ by $4^k$ and multiply $2^k$ back in the end. Find an $A$ satisfying $N-A^2\not\equiv0,4,7 \mod 8$. It's always valid to choose $A=\lfloor\sqrt{N}\rfloor$ or $\lfloor\sqrt{N}\rfloor-1$. (Choose the even one if $N\equiv 1,5 \mod 8$ and the odd one otherwise.) Now $N-A^2$ is an integer in $O(\sqrt{N})$ and you can find $B,C,D$ in $O(\sqrt{N})$ easily.

Edit: Since $B,C,D\le cN^{\frac{1}{4}}$ for some small constant $c$, we can enumerate two of them (suppose $B$ and $C$) and check whether the corresponding $D$ is an integer. This works in $O(\sqrt{N})$. There might be some faster solutions.

StackExchange Computer Science Q#68501, answer score: 4