In this post, I am going to discuss an algorithm for finding the smallest $x$, such that $x^2 \mod K = 0$. This is actually a sub problem of leetcode question 2949. Count Beautiful Substrings II.

This problem is quite trivial if $K$ is prime, and the answer in such case is simply $K$ itself. So how can we find $x$ if $K$ is not prime?

Prime Factorization

First, we need to know that any integer $x$ can be factorized as the product of several prime numbers, i.e.,
$$
x = \prod_i (p_i)^{n_i}
$$
where $p_i$ is the $i^{th}$ prime factor and $n_i$ is the number of times $p_i$ appears in the factorization. What’s more, we can easily get the prime factorization of $x^2$ as
$$
x^2 = \prod_i (p_i)^{ 2 n_i}
$$

The simplest algorithm for prime factorization is shown as follows, and you can check advance algorithms here. We only need to keep dividing the input number with current prime factor $p$ and increase the factor $p$ step by step until we reach the stop condition.

def prime_fac(v):
    p,ans = 2,[]
    while p*p <= v:
        c = 0
        while v % p == 0:
            v = v//p
            c += 1
        if c:
            ans.append( (p,c) )
        p += 1

    if v > 1:
        ans.append( (v, 1) )
    return ans

Note that in the second while loop, if the condition $v \mod p = 0$ is true, then $p$ must be prime.

Otherwise, $p$ must be able to be factorized as the product of smaller numbers $a,b, …$ and $v \mod a =0$ holds as well. However, for these smaller factors, we already divided them in previous while loops, which guarantees the current $v$ cannot modulo them. A contradiction.

The time complexity of the algorithm is $O(\sqrt{v})$ in the worse case when the input is a prime number and we have try to increase $p$ until we reach $\sqrt v$.

Solution

Now we know the prime factorization and how to do it, we can solve our original problem as follows. Since we have $x^2 \mod K = 0$, if we factorize $K$ as
$$
K = \prod_i (q_i)^{m_i}
$$
then all of the prime factors $q_i$ must appears in the prime factorization of $x^2$ at least $m_i$ times.

Therefore, to construct such an $x$ and make it as small as possible, we let $x$ have the same factor as $K$ and we don’t add any extra factors. In other words, $p_i = q_i$. In addition, we also need to guarantee that $2n_i \ge m_i$, and thus the minimum $n_i$ we can have is $\lceil \frac{m_i}{2} \rceil$.

Based on above discussion, the final answer is
$$
x = \prod_i (q_i) ^ {\lceil \frac{m_i}{2} \rceil}
$$

It is also worth noting that for any other $y^2 \mod K = 0$, $y$ must be the multiples of $x$, i.e., $y = nx$ for some $n$. This is because the factorization of $y$ must include all terms from the factorization of $x$ such that $y^2 \mod K = 0$. Therefore, $y \mod x =0$, which gives us the result.

Prime Factorization
Solution