이산 제곱근을 구해보자!
때는 바야흐로 2025년 6월 11일. 일반화학 시험을 공부를 안 해간 나는 문제를 더 풀 수가 없었다. 심지어 조기 퇴실도 안 되는 시험이라 뭐 할지 고민하다가…
이산 제곱근 알고리즘을 구현해보기로 했다!
…만, 당연히 시험 중이라 인터넷 검색도 안 되고 하니 그냥 직접 이산 제곱근을 구하는 알고리즘을 만들어보기로 했다.
이산 제곱근
말 그대로다. $x^2 \equiv d \pmod{p}$를 만족하는 $x$를 $d$의 이산 제곱근이라고 한다. 당연히 2개씩 ($\pm$) 쌍을 지어 나오고, 따라서 $\Bbb{Z}/p\Bbb{Z}$ 위에서 0을 제외하고 절반의 원소만이 이산 제곱근을 가진다. 어떤 원소가 이산 제곱근을 가지는지는 어떻게 알 수 있을까? $a^{\frac{p-1}{2}} \equiv 1 \pmod{p}$라면 $a$는 이산 제곱근을 가지고, $a^{\frac{p-1}{2}} \equiv -1 \pmod{p}$이라면 $a$는 이산 제곱근을 가지지 않는다. (증명은 생략한다. 원시근을 쓰면 유도할 수 있다.) 저 값이 너무 중요한 나머지, $\big(\frac{a}{p}\big) \equiv a^{\frac{p-1}{2}} \pmod{p}$라는 기호도 붙었다. 이걸 르장드르 기호라고 한다. 르장드르 기호 자체에도 여러 중요한 성질들이 있지만, 여기서는 중요하지 않으니 넘어가자.
특정한 소수의 경우에는 이산 제곱근이 매우 쉽게 구해진다. 예를 들어, $p \equiv 3 \pmod{4}$를 만족하는 소수는, $x = d^{\frac{p+1}{4}}$의 꼴로 이산 제곱근을 구할 수 있다. ($x \equiv 1 \pmod{4}$라면 $\frac{p+1}{4}$가 정수가 아니라서 안 된다.) 이게 되는 이유는 페르마의 소정리 때문인데, $a^{p-1} \equiv 1 \pmod{p}$가 $p \nmid a$인 모든 $a$에 대해 성립한다. $(d^\frac{p+1}{4})^2 = d^\frac{p+1}{2} = (d^\frac{p-1}{2}) \cdot d = \big(\frac{d}{p}\big) \cdot d$이고, $d$가 이산 제곱근을 가진다면 $\big(\frac{d}{p}\big) = 1$이 된다.
$p \equiv 1 \pmod{4}$라면 대신 $-1$의 이산 제곱근은 쉽게 구할 수 있다. ($p \equiv 3 \pmod{4}$라면 $(-1)^\frac{p-1}{2} \equiv -1 \pmod{p}$가 되어 이산 제곱근을 가지지 않는다.) 먼저 임의의 $\Bbb{Z}/p\Bbb{Z}$ 위에서 제곱근을 가지지 않는 수 $b$를 하나 찾아오자. 그러면 $\big(\frac{b}{p}\big) \equiv b^\frac{p-1}{2} \equiv -1 \pmod{p}$가 되는데, 이번에는 $\frac{p-1}{4}$가 정수이므로 $b^\frac{p-1}{4}$를 구하면 $-1$의 이산 제곱근을 구할 수 있다.
이제 궁금한 건, $p \equiv 1 \pmod{4}$인 경우에 임의의 $d$에 대한 이산 제곱근을 구하는 것이다. 다시 첫 번째 경우를 살펴보면, $\frac{p+1}{4}$ 말고 다른 분수를 지수에 올려서 저런 형태를 만들 수 있지 않을까? 하는 생각이 든다. $d^{p-1} = 1$이므로 지수에 $\frac{a(p-1)+b}{2b}$ 형태가 오게 하면 될 것 같다. 예를 들어서, $p \equiv 5 \pmod{8}$이면 $\frac{p-1+4}{2 \cdot 4} = \frac{p+3}{8}$이 정수가 된다. 그럼 $d^{\frac{p+3}{8}}$이 $d$의 이산 제곱근이 될까?
\[(d^\frac{p+3}{8})^2 = d^\frac{p+3}{4} = d^{\frac{p-1}{4}} \cdot d\]이번에는 앞에 $d^{\frac{p-1}{2}}$가 아니라 $d^{\frac{p-1}{4}}$가 왔다. $d^{\frac{p-1}{2}} = 1$이니까, $d^{\frac{p-1}{4}} = \pm 1$ 이 성립하기는 할 것이다. 따라서 $d^\frac{p+3}{8}$은 $d$ 또는 $-d$의 이산 제곱근이 된다. 사실 우리는 $-1$의 이산 제곱근을 구하는 법을 이미 알고 있기 때문에, 이를 $i$라 두고 $d^\frac{p+3}{8}$이 $-d$의 이산 제곱근이라면 $d^\frac{p+3}{8}i$가 $d$의 이산 제곱근이 된다는 것은 알 수 있다. 따라서 $p \equiv 5 \pmod{8}$일 때도 이산 제곱근을 구할 수 있게 되었다!
$\frac{p+1}{4}$, $\frac{p+3}{8}$을 일반화해서 $\frac{p+2^k-1}{2^{k+1}}$로 나타낼 수 있을까? $p \equiv 2^k+1 \pmod{2^{k+1}}$이라면 $\frac{p+2^k-1}{2^{k+1}}$가 항상 정수가 되기는 한다. 그리고 모든 홀수 소수는 $a \cdot 2^k + 1$ 꼴로 나타낼 수 있으므로 $p \equiv 2^k+1 \pmod{2^{k+1}}$ 꼴의 소수에 대해 이산 제곱근을 구할 수 있다면 모든 소수에 대해서도 구할 수 있는 셈이 된다. 다만 방금의 $k = 2$인 경우에서 보았듯이, $d^\frac{p+2^k-1}{2^{k+1}}$은 $d$의 이산 제곱근이 아니라 $d^{\frac{p-1}{2^k}} \cdot d$의 이산 제곱근을 구해준다는 것이 문제라는 걸 바로 떠올릴 수 있다. 이러한 ‘오차’ $d^{\frac{p-1}{2^k}} = m$이라 두자. 우리는 $md$의 이산 제곱근을 구할 수 있으므로, $m$의 이산 제곱근을 구할 수 있다면 $d$의 이산 제곱근도 구할 수 있다.
$\big(\frac{d}{p}\big) = d^\frac{p-1}{2} = 1$이므로 $m^{2^{k-1}} = (d^{\frac{p-1}{2^k}})^{2^{k-1}} = d^\frac{p-1}{2} = 1$이다. 따라서 $m$은 $1$의 $2^{k-1}$제곱근 중 하나라는 걸 알 수 있다. 따라서 $\omega’ = -1$의 $2^{k-1}$제곱근, $\omega = (\omega’)^2$이라고 두면 $m = \omega^e = (\omega’^e)^2$ 꼴로 나타내어질 수 있다. 이러한 $e$를 구하기만 하면 끝이다! 길이 보인다. $\omega’$는 매우 쉽게 구할 수 있는데, 아까 $-1$의 이산 제곱근 구했던 것을 응용해서 $\omega’ \equiv b^\frac{p-1}{2^k} \pmod{p}$라고 할 수 있다. 아까 $p \equiv 2^k + 1 \pmod{2^{k+1}}$이라고 가정했으므로 $\omega’$는 잘 구해진다.
이제 남은 건 $e$를 구하는 것이다. $1$의 $2^{k-1}$ 제곱근은 진짜 $2^{k-1}$개 있어서 저걸 다 순회하면 시간이 매우 오래 걸린다. 분할 정복 거듭제곱과 같은 아이디어를 사용하자. $e = \displaystyle\sum_{i=0}^{k-2} e_i 2^i$라고 두면 $m = \omega^e = \displaystyle\prod_{i=0}^{k-2} \omega^{e_i 2^i}$이다. $\omega^{2^{k-2}} = -1$이므로 이걸 통채로 $2^{k-2}$제곱을 하면 $m^{2^{k-2}} = \displaystyle\prod_{i=0}^{k-2} (-1)^{e_i 2^i} = (-1)^{e_0}$이 되고, 따라서 이 값이 $1$이라면 $e_0 = 0$, $-1$이라면 $e_0 = 1$이 되는 걸 알 수 있다. $e_0$을 구했으니 $m_1 = m \cdot \omega^{-e_0}$로 두고, $s_0 = {\omega’}^{e_0}$이라 두자. 이번에는 $2^{k-3}$제곱을 해서 $e_1$을 구하고, $m_2 = m_1 \cdot \omega^{-2e_1}$ / $s_1 = s_0 \cdot \omega’^{2e_1}$로 두자. 이걸 $k-1$번 반복하면 $s^2 = m$인 $s$를 구할 수 있다. 그러므로 $(d^\frac{p+2^k-1}{2^{k+1}} \cdot s^{-1})^2 \equiv md \cdot m^{-1} \equiv d \pmod{p}$가 되어, $d$의 이산 제곱근을 구했다!
약간만 더 최적화를 해 보자. $s$를 구할 때 $\omega^{-e_i2^i}$를 계속 곱셈 역원으로 구하면서 구하지 말고, 처음부터 $\omega’$ 대신 ${\omega’}^{-1} = \bar\omega’$를 사용하고, $\bar s_1 = \bar\omega^{e_0}$, $\bar s_n = \bar s_{n-1} \cdot \bar\omega^{e_n2^n}$을 구하자. 그럼 결과적으로 $\bar s = s^{-1}$ 을 바로 구할 수 있다. 또한 $\bar\omega’$ 역시 $-1$의 $2^{k-1}$이므로, $\bar\omega’$ 대신 그냥 $\omega’$를 사용해도 된다.
$k = {\cal O}(\log p)$이고, $k$번 순회하는 과정에서 $2^{k}$로 분할 정복 거듭제곱을 하므로 ${\cal O}(k + \log p)$만큼이 더 든다. 따라서 최종 시간복잡도는 ${\cal O}(\log^2 p)$.
아래는 Python 구현이다. $p \equiv 3 \pmod{4}$와 $p \equiv 5 \pmod{8}$인 경우도 $p \equiv 2^k+1 \pmod{2^{k+1}}$을 만족한다 할 수 있으므로 하나의 코드로 충분하다. $p=2$일 때만 따로 처리해 주자.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def discrete_sqrt(n, p):
if p == 2: return n%2
k = 0
while (p-1) % (2 << k) == 0: k += 1
t = 1 << k
w0 = 1
while pow(w0, (p-1)//2, p) == 1: w0 += 1
w0 = pow(w0, (p-1)//t, p)
m = pow(n, (p-1)//t, p)
s = 1
for i in range(k-2, -1, -1):
if pow(m, 1 << i, p) != 1:
s = s*w0 % p
m = m*w0*w0 % p
w0 = w0*w0 % p
return s * pow(n, (p+t-1)//(2*t), p) % p