I participated in TetCTF as a member of Super Guesser. 

I was first to [solve all three crypto challenges]. Our team finished 2nd place :)


unimplemented (1st solve)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
from collections import namedtuple
from Crypto.Util.number import getPrime
import random
 
Complex = namedtuple("Complex", ["re""im"])
 
 
def complex_mult(c1, c2, modulus):
    return Complex(
        (c1.re * c2.re - c1.im * c2.im) % modulus,  # real part
        (c1.re * c2.im + c1.im * c2.re) % modulus,  # image part
    )
 
 
def complex_pow(c, exp, modulus):
    result = Complex(10)
    while exp > 0:
        if exp & 1:
            result = complex_mult(result, c, modulus)
        c = complex_mult(c, c, modulus)
        exp >>= 1
    return result
 
 
def generate_key_pair(nbits):
    while True:
        p = getPrime((nbits + 3// 4)
        q = getPrime((nbits + 3// 4)
        n = (p ** 2* (q ** 2)
        if n.bit_length() == nbits:
            return (p, q), n
 
 
def pad(data, length):
    assert len(data) < length
    pad_length = length - len(data) - 1
    pad_data = bytes(random.choices(range(1256), k=pad_length))
    sep = b'\x00'
    return pad_data + sep + data
 
 
def unpad(data):
    assert b"\x00" in data, "incorrect padding"
    return data.split(b"\x00"1)[1]
 
 
def encrypt(public_key, plaintext):
    n = public_key
    plaintext = pad(plaintext, 2 * ((n.bit_length() - 1// 8))
    m = Complex(
        int.from_bytes(plaintext[:len(plaintext) // 2], "big"),
        int.from_bytes(plaintext[len(plaintext) // 2:], "big")
    )
    c = complex_pow(m, 65537, n)
    return (c.re.to_bytes((n.bit_length() + 7// 8"big")
            + c.im.to_bytes((n.bit_length() + 7// 8"big"))
 
 
def decrypt(private_key, ciphertext):
    # TODO
    raise Exception("unimplemented")
 
 
def main():
    private_key, public_key = generate_key_pair(2021)
    from secret import flag
 
    print("private_key =", private_key)
    print("public_key =", public_key)
    print("ciphertext =", encrypt(public_key, flag))
 
 
if __name__ == '__main__':
    main()
 
# Output:
# private_key = (128329415694646850105527417663220454989310213490980740842294900866469518550360977403743209328130516433033852724185671092403884337579882897537139175073013,
#                119773890850600188123646882522766760423725010264224559311769920026142724028924588464361802945459257237815241227422748585976629359167921441645714382651911)
# public_key = 236252683050532196983825794701514768601125614979892312308283919527619033977486749228418695923608569040825653688303374445536392159719426640289893369552258923597180869981053519695297428186215135878525530974780390951763007339139013157234202093279764459949020588291928614938201565110828675907781512603972957429701280916745719458738970910789383870206038035515777549907045905872280803964436193687072794878040018900969772972761081589529671158140590712503582004892067155769362463889653489918914397872964087471457070748108694165412065471040954221707557816986272311750297566993468288899523479556822418109112211944932649
# ciphertext = b'\x00h\xbe\x94\x8c\xcd\xdd\x04\x80\xf4\x9d\t\xd8\x8dO\x08\xf1\xd1\xc4\xb9\xa06\xe7\xe3\xb6\xc3\x01+\xa9\xf2\xb9\xe8\x8d\xe6\xc9\x0c_#\x93\x11\xad\x0f\x90\xd3\x0b6\xb0n\x13\xea~"V6\xdbA&\x87\xfe\xa3C\xcb\x16\xae\xd9\x83\xdbU\xc6\x06\xcd\x9a\x94\xa9\xce\x15{d\x95s\xc2\xfb\x90q\xe7\x02\xa2\x081:_C\xc68\x00y\x7fj4@\xd2\xcdE\x06\x943\xbe\xbcC3\xca\x91\xb4\x0e}C\xab\xff?X\xc30u\x069:Dc\xb5\xdc\x9b0%\x98\xbd\xd9\x13\xc0\x02w\xc5\xe5:\xca\xcf\x0c\xab\xc2\x9b}\xab\xd0\xcc\xbc\x0f\x9e9\t\xf7M\xb3\xed\x86\xb5E\x8b\xbc4\xfaH\x9b4\x1c\xc4\xab\xc0\xaf\x8a5XcX\x19K\xed\x19\xe1\x1c\xd0\x1e\x97c\x9fF:L\x9d\x90p\x99\xb8\xde\xcblM\xb3\x80sSL\xe1\xa4\xd6,\x81\xd6\x9c\xf1\xbb\x9c)\xf03\x155\xc0\xd5v\x13\xd6#\xb7\x19\xdea%\xce<\x1c\xf7\xf2!;\xc1\xd7w\xd1\xc8\x8d/\xaew\xa1\x1d\xc5(\xc8\x9d\x82v\xf6j\x90A,e\xbd\xa7]\x10\x8f\xe5\xe7\x93}:\xdf1~\xec\xd0-o`\r\x96\xe4\x03\xb9E\x9fdF\xc3\xf8L\xa0\xda\xf0E[\xf7\x02\xef|*\x08\x8a5pI&\xa9i\x0fZ\xa8\xb3H\xed\xe8v\xc4H\xff\xdb\xcb\x00\xf1\xae\x9bO\x18\xd5\xd8&p\xf5\xf6\xe9\xf1\x1brs\xc2\r\xceZ\xd0\xb24\x97+\x98b\x0e\xbb\xb8.\x8dT\xe4"\xad\xe4\xa3f\xd0M\xbf\xafX\xbb"[\x99\xdap\xa5\xcfT2Wx\x87M\x7f\x99!>B[Q\x04\xf6\x03\xbc\x84\xf4\xdfj\xdd1^I\x1a\x05\x81\x91\xde9Mf*\x8e\x8d\xe64\xf8\x93\x99&yP\xcd\x00!\x82\xab\xbcy\xed\xf1\x13\xd3\x81\xeaz\xbbP>\x9a2\x8c\x08\x0es\xbc\xa9\xf6\xa3\x8c\xb0\xb9t\xd9?\x06@\xc9\x90\xb7\xa7<\x85\xeb\x1a\x88#\x1c\xc3 \xec\xc7\x94d\x99\xd6\x8e>\x06\xf8Y\xf4\x19\xcaI/hy\x18\x8e\x0e8\xf8\r\xbb\xf6\x11\xb9\x8dCWB6 '
 
cs


Introduction

We are given all the key values and the ciphertext, so we just need to work on the decryption algorithm.

Our encryption is a very RSA-like algorithm, with multiplication defined as complex multiplication, with the Re/Im values taken modulo $n$.


Strategy

Let's think about how RSA decryption works first, and start from there. Denote $n=pq$ with $p, q$ primes.

We know that $\text{gcd}(a, p) = 1$ implies $a^{p-1} \equiv 1 \pmod{p}$. We know that $\text{gcd}(a, q) = 1$ implies $a^{q-1} \equiv 1 \pmod{q}$.

Therefore, we see that with $\text{gcd}(a, p)=1$ and $\text{gcd}(a, q)=1$, we have $a^{(p-1)(q-1)} \equiv 1 \pmod{pq}$.

This is because $a^{(p-1)(q-1)}$ must be $1$ modulo each $p, q$, so we may apply Chinese Remainder Theorem.

Now, with the encryption exponent $e$, we can get decryption exponent by taking modulo inverse mod $(p-1)(q-1)$.


This problem should not be so different. We'll denote $a \equiv b \pmod{n}$ if $a, b$ has same Re/Im values $\pmod{n}$.

We should find a function $f$ that guarantees $a^{f(p)} \equiv 1 \pmod{p^2}$ if $\text{gcd}(a, p) = 1$, and $a^{f(q)} \equiv 1 \pmod{q^2}$ if $\text{gcd}(a, q) = 1$.

  • Immediate Question : What does $\text{gcd}$ mean in this setting? I'll explain this later.

Anyways, if we find such $f$, we know $a^{f(p)f(q)} \equiv 1 \pmod{p^2q^2}$ must hold for the same reason.

We can then find the modular inverse of $e = 65537$ mod $f(p)f(q)$ and get the decryption exponent.


For now, let's just hope that we have $\text{gcd}(a, p) = \text{gcd}(a, q) = 1$ for our ciphertext. (whatever that gcd means)

Now we have to find what the function $f$ is all about. There are at least 5 ways to do this. (I used Method 2.)


Finish

  • Method 1. We simply read the code of unevaluated.
  • It's easy to see that "unevaluated" challenge gives us some discrete logarithm problem.
  • If we read it more carefully, we see that $g = (re, im)^{24}$ precisely has order $pqr = p(p-1)(p+1)/24$. 
  • Therefore, we can see that $(re, im)$ has an order that divides $p(p-1)(p+1)$.
  • We can now guess that $f(p) = p^3 - p$ should work, and it indeed does. Done :)
  • Method 2. We run some tests for small $p$, and guess our way to find $f$.
  • Let's just fix some small primes $p$ and see when $a^{???} \equiv 1 \pmod{p^2}$. 
  • Let's look at the minimum exponent that satisfies the modulo equation. It can be seen that it is a divisor of $p(p-1)(p+1)$.
  • Therefore, we can now guess that $f(p) = p^3 - p$ should work, and it indeed does. Done :)
  • Method 3. We actually do some mathematics to find our answer.
  • We define $\text{gcd}(a+bi, p)$ as $\text{gcd}(a^2+ b^2, p)$. This is highly motivated by the norm map.
  • We first claim that for odd primes $p$ and $\text{gcd}(a+bi, p)=1$, $(a+bi)^{p^2-1} \equiv 1 \pmod{p}$. 
  • Case 1. $p \equiv 3 \pmod{4}$.
  • This is the easy case, as our complex numbers can be now viewed as elements of $\mathbb{F}_{p^2}$.
  • Since the multiplicative group of this extension field has $p^2-1$ elements, we are now done.
  • Case 2. $p \equiv 1 \pmod{4}$.
  • We prove that $(a+bi)^{p-1} \equiv 1 \pmod{p}$. This is enough to prove our result.
  • Denote $(a+bi)^{p-1} \equiv c + di \pmod{p}$, and note that $(a-bi)^{p-1} \equiv c - di \pmod{p}$ by taking conjugates.
  • Now we take an $x$ such that $x^2 \equiv -1 \pmod{p}$. We can now substitute $i$ as $x$ without an issue.
  • This shows $1 \equiv (a+bx)^{p-1} \equiv c+ dx \pmod{p}$, and $1 \equiv (a-bx)^{p-1} \equiv c - dx \pmod{p}$.
  • Combining these two equations directly gives us $c = 1$, and $d = 0$. Note that $x$ is nonzero.
  • Note that $a+bx$ being nonzero $\pmod{p}$ is equivalent to $\text{gcd}(a^2 + b^2, p) = 1$. 
  • Now we prove the main result : $(a+bi)^{p^3 - p} \equiv 1 \pmod{p^2}$ if $\text{gcd}(a+bi, p) = 1$.
  • We know that $(a+bi)^{p^2-1} \equiv 1 \pmod{p}$, so we can write $(a+bi)^{p^2-1} \equiv 1 + pc \pmod{p^2}$ for some complex number $c$.
  • Now we directly use Binomial Theorem to get $(a+bi)^{p^3-p} \equiv (1+pc)^p \equiv 1 \pmod{p^2}$. Done.
  • This shows that we can use $f(p) = p^3-p$, with proof. It should work. Done :)
  • Method 4. Google with keywords "Gaussian Integers" and "RSA"
  • This leads to various resources, such as https://www.diva-portal.org/smash/get/diva2:1170568/FULLTEXT01.pdf
  • Method 5. Lagrange's Theorem
  • We start by defining $\text{gcd}(a+bi, p) = \text{gcd}(a^2 +b^2, p)$, just like in Method 3.
  • It's not hard to prove that the elements with $\text{gcd}(a^2+b^2, p)=1$ form a multiplicative group. 
  • Here, we look at elements $a+bi$ in $\pmod{p^2}$. Therefore, there are at most $p^4$ elements.
  • If we count the number of elements with $\text{gcd}(a^2 + b^2, p)=1$, we can use that as $f(p)$, by Lagrange's Theorem.
  • In other words, we use that $a^{|G|} = e$ for all finite groups $G$, identity element $e$, and element $a \in G$. 
  • Case 1. $p \equiv 3 \pmod{4}$.
  • For $\text{gcd}(a^2+b^2,p) \neq 1$, we need both $a, b$ to be multiples of $p$. There are $p^2$ such elements.
  • Therefore, we may take $f(p) = p^4 - p^2$ in this case. Done!
  • Case 2. $p \equiv 1 \pmod{4}$.
  • We only need to look at $0 \le a, b < p$ and multiply by $p^2$ in the end.
  • If $a=0$, only $b=0$ satisfies $\text{gcd}(a^2+b^2,p) \neq 1$. If $1 \le a < p$, there are $2$ such $b$s. 
  • Therefore, there are $p^2-1-2(p-1) = (p-1)^2$ solutions. Multiplying by $p^2$, we have $f(p) = p^2(p-1)^2$. 
  • Since casework is boring, we just take the LCM of the two expressions and select $f(p) = p^2(p-1)^2(p+1)$. Done.
  • Looking at the flag, the mentions of "counting" seems to imply that this method is the intended solution.

I wish (at least) Method 1 was blocked by allowing contestants to see "unevaluated" only after solving this challenge. :(

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
private_key = (128329415694646850105527417663220454989310213490980740842294900866469518550360977403743209328130516433033852724185671092403884337579882897537139175073013,
               119773890850600188123646882522766760423725010264224559311769920026142724028924588464361802945459257237815241227422748585976629359167921441645714382651911)
public_key = 236252683050532196983825794701514768601125614979892312308283919527619033977486749228418695923608569040825653688303374445536392159719426640289893369552258923597180869981053519695297428186215135878525530974780390951763007339139013157234202093279764459949020588291928614938201565110828675907781512603972957429701280916745719458738970910789383870206038035515777549907045905872280803964436193687072794878040018900969772972761081589529671158140590712503582004892067155769362463889653489918914397872964087471457070748108694165412065471040954221707557816986272311750297566993468288899523479556822418109112211944932649
ciphertext = b'\x00h\xbe\x94\x8c\xcd\xdd\x04\x80\xf4\x9d\t\xd8\x8dO\x08\xf1\xd1\xc4\xb9\xa06\xe7\xe3\xb6\xc3\x01+\xa9\xf2\xb9\xe8\x8d\xe6\xc9\x0c_#\x93\x11\xad\x0f\x90\xd3\x0b6\xb0n\x13\xea~"V6\xdbA&\x87\xfe\xa3C\xcb\x16\xae\xd9\x83\xdbU\xc6\x06\xcd\x9a\x94\xa9\xce\x15{d\x95s\xc2\xfb\x90q\xe7\x02\xa2\x081:_C\xc68\x00y\x7fj4@\xd2\xcdE\x06\x943\xbe\xbcC3\xca\x91\xb4\x0e}C\xab\xff?X\xc30u\x069:Dc\xb5\xdc\x9b0%\x98\xbd\xd9\x13\xc0\x02w\xc5\xe5:\xca\xcf\x0c\xab\xc2\x9b}\xab\xd0\xcc\xbc\x0f\x9e9\t\xf7M\xb3\xed\x86\xb5E\x8b\xbc4\xfaH\x9b4\x1c\xc4\xab\xc0\xaf\x8a5XcX\x19K\xed\x19\xe1\x1c\xd0\x1e\x97c\x9fF:L\x9d\x90p\x99\xb8\xde\xcblM\xb3\x80sSL\xe1\xa4\xd6,\x81\xd6\x9c\xf1\xbb\x9c)\xf03\x155\xc0\xd5v\x13\xd6#\xb7\x19\xdea%\xce<\x1c\xf7\xf2!;\xc1\xd7w\xd1\xc8\x8d/\xaew\xa1\x1d\xc5(\xc8\x9d\x82v\xf6j\x90A,e\xbd\xa7]\x10\x8f\xe5\xe7\x93}:\xdf1~\xec\xd0-o`\r\x96\xe4\x03\xb9E\x9fdF\xc3\xf8L\xa0\xda\xf0E[\xf7\x02\xef|*\x08\x8a5pI&\xa9i\x0fZ\xa8\xb3H\xed\xe8v\xc4H\xff\xdb\xcb\x00\xf1\xae\x9bO\x18\xd5\xd8&p\xf5\xf6\xe9\xf1\x1brs\xc2\r\xceZ\xd0\xb24\x97+\x98b\x0e\xbb\xb8.\x8dT\xe4"\xad\xe4\xa3f\xd0M\xbf\xafX\xbb"[\x99\xdap\xa5\xcfT2Wx\x87M\x7f\x99!>B[Q\x04\xf6\x03\xbc\x84\xf4\xdfj\xdd1^I\x1a\x05\x81\x91\xde9Mf*\x8e\x8d\xe64\xf8\x93\x99&yP\xcd\x00!\x82\xab\xbcy\xed\xf1\x13\xd3\x81\xeaz\xbbP>\x9a2\x8c\x08\x0es\xbc\xa9\xf6\xa3\x8c\xb0\xb9t\xd9?\x06@\xc9\x90\xb7\xa7<\x85\xeb\x1a\x88#\x1c\xc3 \xec\xc7\x94d\x99\xd6\x8e>\x06\xf8Y\xf4\x19\xcaI/hy\x18\x8e\x0e8\xf8\r\xbb\xf6\x11\xb9\x8dCWB6 '
 
= private_key[0]
= private_key[1]
= public_key
= 65537
 
= (p ** 3 - p) * (q ** 3 - q)
= inverse(e, v) 
 
re = bytes_to_long(ciphertext[ : len(ciphertext) // 2])
im = bytes_to_long(ciphertext[len(ciphertext) // 2 : ])
 
ctxt = Complex(re, im)
 
ptxt = complex_pow(ctxt, d, n)
 
print(long_to_bytes(ptxt.re) + long_to_bytes(ptxt.im))
 
# TetCTF{c0unt1ng_1s_n0t_4lw4ys_34sy-vina:*100*48012023578024#}
cs


uncollidable (4th solve)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import json
import os
from hashlib import sha256
from typing import Callable, Dict
 
H: Callable[[bytes], bytes] = lambda s: sha256(s).digest()
 
 
def get_mac(key: bytes, data: bytes) -> bytes:
    """Just a MAC generation function similar to HMAC.
    Reference link: https://en.wikipedia.org/wiki/HMAC#Implementation
    """
    if len(key) >= 32:
        key = H(key)
 
    if len(key) < 32:
        key += b"\x00" * (32 - len(key))
 
    inner_pad = bytearray(32)
    outer_pad = bytearray(32)
    for i in range(32):
        inner_pad[i] = key[i] ^ 0x20
        outer_pad[i] = key[i] ^ 0x21
 
    return H(outer_pad + H(inner_pad + data))
 
 
class SecureStorage:
    """Changes on data stored in this storage can be detected :) """
 
    def __init__(self, debug=False):
        self.keys: Dict[str, bytes] = {}  # key ID -> key
        self.data: Dict[bytes, bytes] = {}  # MAC -> data
        self.debug = debug
        if debug:
            # when debug mode is on, unique arguments to `store_data` will be
            # logged.
            self.store_data_args = set()
 
    def generate_key(self, key_id: str-> None:
        """Generate a new key."""
        if key_id in self.keys:
            raise Exception("duplicated key id")
        self.keys[key_id] = os.urandom(32)
 
    def import_key(self, key_id: str, key: bytes) -> None:
        """Import an external key."""
        if key_id in self.keys:
            raise Exception("duplicated key id")
        self.keys[key_id] = key
 
    def store_data(self, key_id: str, data: bytes) -> bytes:
        """Store data and return a MAC that can be used to retrieve the data
        later."""
        if key_id not in self.keys:
            raise Exception("key not found")
 
        key = self.keys[key_id]
        if self.debug:
            self.store_data_args.add((key, data))
 
        mac = get_mac(key, data)
        if mac in self.data:
            raise Exception("data already stored")
 
        self.data[mac] = data
        return mac
 
    def retrieve_data(self, key_id: str, mac: bytes) -> bytes:
        """Retrieve data previously stored with `store_data`."""
        if key_id not in self.keys:
            raise Exception("key not found")
        if mac not in self.data:
            raise Exception("data not found")
 
        key = self.keys[key_id]
        data = self.data[mac]
        if get_mac(key, data) != mac:
            raise Exception("data has been tampered")
        return data
 
    def report_bug(self-> int:
        """Claim your bounty here :) """
        if self.debug:
            # check for a massive collision bug
            if len(self.store_data_args) >= 10 and len(self.data) == 1:
                # check if collisions happened through different key sizes
                key_lengths = [len(key) for (key, _) in self.store_data_args]
                if min(key_lengths) < 32 <= max(key_lengths):
                    # Congrats :)
                    from secret import flag
                    return int.from_bytes(flag, "big")
        return 0
 
 
def main():
    ss = SecureStorage(debug=True)
    for _ in range(100):
        try:
            request = json.loads(input())
 
            if request["action"== "generate_key":
                key_id = request["key_id"]
                ss.generate_key(key_id)
                print(json.dumps({}))
 
            if request["action"== "import_key":
                key_id = request["key_id"]
                key = bytes.fromhex(request["key"])
                ss.import_key(key_id, key)
                print(json.dumps({}))
 
            elif request["action"== "store_data":
                key_id = request["key_id"]
                data = bytes.fromhex(request["data"])
                mac = ss.store_data(key_id, data)
                print(json.dumps({"mac": mac.hex()}))
 
            elif request["action"== "retrieve_data":
                key_id = request["key_id"]
                mac = bytes.fromhex(request["mac"])
                data = ss.retrieve_data(key_id, mac)
                print(json.dumps({"data": data.hex()}))
 
            elif request["action"== "report_bug":
                print(json.dumps({"bounty": ss.report_bug()}))
 
        except EOFError:
            break
        except Exception as err:
            print(json.dumps({"error"str(err)}))
 
 
if __name__ == '__main__':
    main()
 
cs


Introduction

We basically need a big collision of the MAC - 10 (key, data) values that map to a single MAC value.

We also need to use at least one key with length less than 32, and at least one key with length no less than 32 as well. 


Strategy & Finish

Let's actually look at the MAC code. We see that some preparation are done to make the key exactly 32 bytes.

If the length is at least 32 bytes, we apply SHA256. If not, we append some zeros at the end to make it 32 bytes.

We immediately get some collision cooking - a key full of zeros will lead to a same MAC if the key length is less than 32 bytes.

Of course, this is easily generalizable - if the "final" 32 byte key ends with a lot of zeros, we get a lot of collisions!


This is extremely suspicious - we just need a key that is at least 32 bytes to go along with it.

This means we need a 32 byte value with SHA256 ending with a lot of zeros - 9 bytes of it, to be exact.


Now this is a problem that I have seen before - DragonCTF's bitflip2. SHA256 ending with a lot of zeros leads us to bitcoin PoW.

We can find a bunch of values with SHA256 ending with 8 bytes worth of zeros in bitcoin. 

Luckily, the value we (Perfect Guesser) used at the time to solve bitflip2 already had 9 bytes of zeros in the end in the SHA256.


This solves the problem. This was a very good review exercise for me :)


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
HOST = "139.162.5.141"
PORT = 5555
 
= remote(HOST, PORT)
 
key_id = ["0""1""2""3""4""5""6""7""8""9"]
 
print(long_to_bytes(61342502866914327869718742967850222962760441300896006743258488289386736694996435440480147718715690291226752007898889408599107845232272253))
tt = '294724a63e0fda8c731d9612b0a8e8b9bc0ec087ca9920c8488c5dd1df94ebff'
tt = bytes.fromhex(tt)
 
data = "1234"
 
for i in range(010):
    key = '294724a63e0fda8c731d9612b0a8e8b9bc0ec087ca9920c8488c5dd1df94ebff'
    if i >= 1:
        key = bytes.fromhex(sha256(bytes.fromhex(key)).hexdigest())
        key = key[:-i]
        key = key.hex()
    request = {
        "action""import_key",
        "key_id": key_id[i],
        "key": key
    }
    r.sendline(json.dumps(request))
    print(r.readline())
    request = {
        "action""store_data",
        "key_id": key_id[i],
        "data": data
    }
    r.sendline(json.dumps(request))
    print(r.readline())
 
request = {
    "action""report_bug"
}
 
r.sendline(json.dumps(request))
 
print(r.readline())
 
# TetCTF{HM4C_c4n_b3_m1sus3d-viettel:*100*718395803842748#}
cs


unevaluated (1st solve)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
from collections import namedtuple
from Crypto.Util.number import getPrime, isPrime, getRandomRange
 
Complex = namedtuple("Complex", ["re""im"])
 
 
def complex_mult(c1, c2, modulus):
    return Complex(
        (c1.re * c2.re - c1.im * c2.im) % modulus,  # real part
        (c1.re * c2.im + c1.im * c2.re) % modulus,  # image part
    )
 
 
def complex_pow(c, exp, modulus):
    result = Complex(10)
    while exp > 0:
        if exp & 1:
            result = complex_mult(result, c, modulus)
        c = complex_mult(c, c, modulus)
        exp >>= 1
    return result
 
 
class ComplexDiffieHellman:
    @staticmethod
    def generate_params(prime_length):
        # Warning: this may take some time :)
        while True:
            p = getPrime(prime_length)
            if p % 4 == 3:
                if p % 3 == 2:
                    q = (p - 1// 2
                    r = (p + 1// 12
                    if isPrime(q) and isPrime(r):
                        break
                else:
                    q = (p - 1// 6
                    r = (p + 1// 4
                    if isPrime(q) and isPrime(r):
                        break
        n = p ** 2
        order = p * q * r
        while True:
            re = getRandomRange(1, n)
            im = getRandomRange(1, n)
            g = complex_pow(Complex(re, im), 24, n)
            if (
                    complex_pow(g, order, n) == Complex(10)
                    and complex_pow(g, order // p, n) != Complex(10)
                    and complex_pow(g, order // q, n) != Complex(10)
                    and complex_pow(g, order // r, n) != Complex(10)
            ):
                return g, order, n
 
    def __init__(self, params=None, prime_length=128, debug=False):
        if not debug:
            raise Exception("security unevaluated")
        if params is None:
            params = ComplexDiffieHellman.generate_params(prime_length)
        self.g, self.order, self.n = params
 
    def get_public_key(self, private_key):
        return complex_pow(self.g, private_key, self.n)
 
    def get_shared_secret(self, private_key, other_public_key):
        return complex_pow(other_public_key, private_key, self.n)
 
 
def main():
    from os import urandom
    private_key = urandom(32)
    k = int.from_bytes(private_key, "big")
 
    cdh = ComplexDiffieHellman(debug=True)
    print("g =", cdh.g)
    print("order =", cdh.order)
    print("n =", cdh.n)
    print("public_key =", cdh.get_public_key(k))
 
    # Solve the discrete logarithm problem if you want the flag :)
    from secret import flag
    from Crypto.Cipher import AES
    if len(flag) % 16 != 0:
        flag += b"\x00" * (16 - len(flag) % 16)
    print("encrypted_flag = ",
          AES.new(private_key, AES.MODE_ECB).encrypt(flag))
 
 
if __name__ == "__main__":
    main()
 
# Output:
# g = Complex(re=20878314020629522511110696411629430299663617500650083274468525283663940214962,
#             im=16739915489749335460111660035712237713219278122190661324570170645550234520364)
# order = 364822540633315669941067187619936391080373745485429146147669403317263780363306505857156064209602926535333071909491
# n = 42481052689091692859661163257336968116308378645346086679008747728668973847769
# public_key = Complex(re=11048898386036746197306883207419421777457078734258168057000593553461884996107,
#                      im=34230477038891719323025391618998268890391645779869016241994899690290519616973)
# encrypted_flag = b'\'{\xda\xec\xe9\xa4\xc1b\x96\x9a\x8b\x92\x85\xb6&p\xe6W\x8axC)\xa7\x0f(N\xa1\x0b\x05\x19@<T>L9!\xb7\x9e3\xbc\x99\xf0\x8f\xb3\xacZ:\xb3\x1c\xb9\xb7;\xc7\x8a:\xb7\x10\xbd\x07"\xad\xc5\x84'
 
cs


Introduction

We are given the same setup from "unimplemented". We are solving $g^x \equiv target \pmod{n}$ given $g, target, n$.

We also know that the order of $g$ is $order = p \cdot q \cdot r = p(p-1)(p+1)/24$, where $p, q, r$ are all primes.

We now have to solve a discrete logarithm problem to get a private key, which is used to encrypt our flag using AES.


Strategy 

First, some basic steps. From $n = p^2$, we can directly calculate the value of $p$.

From the parameter generation process, we can also calculate $q, r$. 

I actually computed $p$ by $\text{gcd}(n, order)$, and factorized $qr = order/ p$ using alpetron. 

UPD: This is why my values of $q, r$ in this writeup and solution code is swapped instead of the original $q, r$ given in the code.

Anyway, we now know $p, q, r$ and can focus on our task, the discrete logarithm.

  • Idea 1. The flaw in parameter generation.
  • There is actually a very important flaw in parameter generation. 
  • Note that $p, q, r$ are around 128 bits. Therefore, the order $pqr$ is around $128 \cdot 3 = 384$ bits.
  • However, the private key used in the AES and the discrete logarithm problem is 32 bytes, i.e. $256$ bits.
  • Pohlig-Hellman style thinking shows we can compute $x \pmod{p}$, $x \pmod{q}$, $x \pmod{r}$ separately. 
  • Because $x$ is $256$ bits, we can simply solve two of these three problems and combine using CRT!!
  • Idea 2. The (multiplicative) norm map $N(a+bi) = a^2 + b^2$. 
  • Note : I proved stuff for "unimplemented" (i.e. Method 3) after I solved all three challenges. However, the proof ideas are used here.
  • The norm map allows us to move from weird complex number stuff to values in $\pmod{p}$, where we are more comfortable.
  • Maybe this idea reminds us of MOV attack? I don't know. Anyways, let's finish the problem.

Finish

Let's first look at $x \pmod{r}$. We note that $r|(p-1)$, so norm map will work nicely here.

  • We want $(g^{pq})^x \equiv target^{pq} \pmod{n}$, so $N(g^{pq})^x \equiv N(target^{pq}) \pmod{p}$. 
  • If we calculate $N(g^{pq})$ and $N(target^{pq})$ we get values that are nonzero, and also not equal to $1$.
  • This is expected to happen, since $pq$ is relatively prime to $p-1$ and $g$ has high order. 
  • Therefore, we can solve the discrete logarithm problem in $\mathbb{F}_p$ to get information on $x$.
  • We also see that (with Sage) $N(g^{pq})$ has order $r$ in $\mathbb{F}_p$. Therefore, we get the entire $x \pmod{r}$ here.

Next, we'll look at $x \pmod{p}$. If we get this, we'll solve the problem.

  • We want $(g^{qr})^x \equiv target^{qr} \pmod{n}$. Note that $24qr \equiv 0 \pmod{p^2-1}$. 
  • Here, the $24$ comes from the fact that $g$ is actually defined as $(re, im)^{24}$ for some $re, im$. 
  • This shows that, from proof in Method 3 above, we must have $g^{qr} \equiv target^{qr} \equiv 1 \pmod{p}$.
  • Therefore, we have $N(g^{qr}) \equiv N(target^{qr}) \equiv 1 \pmod{p}$ as well. 
  • However, computation shows us that $N(g^{qr}) , N(target^{qr})$ are not $1 \pmod{p^2}$. This is not unexpected.
  • We can now solve $x \pmod{p}$ with $N(g^{qr})^x \equiv N(target^{qr}) \pmod{p^2}$. 
  • We simply note that $(1+pa)^x \equiv 1 +pax \pmod{p^2}$ using Binomial Theorem.
  • Write $N(g^{qr}) \equiv 1 +pa \pmod{p^2}$ and $N(target^{qr}) \equiv 1 + pb \pmod{p^2}$, where $a, b$ are nonzero $\pmod{p}$.
  • We now have, simply, $ax \equiv b \pmod{p}$. This can be solved with modular inverse.
  • This looks quite like the Paillier Cryptosystem (which I learned by solving zer0ptsCTF Dirty Laundry) so that's cool.

When I first solved this problem, I found most of these insight by experiments. 

I did think of the norm map first, solved $\pmod{r}$ part based on it, and was quite surprised when $N(g^{qr})$ was $1 \pmod{p}$. 

Then I realized that I could use Binomial Theorem to solve the challenge. Looking back on "unimplemented" during write-up, I found out that most of my "surprises" during the solve was not actually surprising facts :) This is a very nice challenge, so thanks to ndh.


UPD: After some discussion with ndh in the CryptoHack server, it turns out I overcomplicated this problem a bit.

We can just solve for $N(g)^x \equiv N(target) \pmod{p^2}$ and solve the discrete logarithm problem here. :P

UPD: After even more discussion, it turns out that the above discrete logarithm in $\pmod{p^2}$ is not that easy.

Refer to the CryptoHack discord, or the writeup by Mystiz here, and TheBlueFlame121 here for more details. Highly suggest you read them.

These two writeups linked above have some important facts about discrete logarithm calculation in Sage in general. :)

In CryptoHack discord #ctf-challenges, $\$$in has a summarization on this subject, and there are even more discussions with smart people!

UPD: CryptoHack discord #ctf-challenges has a very in-depth analysis and discussion on this challenge, a must read.

UPD: CryptoHack has posted a solid review on these challenges. See https://blog.cryptohack.org/tetctf-2021

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def norm(x):
    return x.re * x.re + x.im * x.im
 
 
= Complex(re=20878314020629522511110696411629430299663617500650083274468525283663940214962,
            im=16739915489749335460111660035712237713219278122190661324570170645550234520364)
order = 364822540633315669941067187619936391080373745485429146147669403317263780363306505857156064209602926535333071909491
= 42481052689091692859661163257336968116308378645346086679008747728668973847769
public_key = Complex(re=11048898386036746197306883207419421777457078734258168057000593553461884996107,
                     im=34230477038891719323025391618998268890391645779869016241994899690290519616973)
encrypted_flag = b'\'{\xda\xec\xe9\xa4\xc1b\x96\x9a\x8b\x92\x85\xb6&p\xe6W\x8axC)\xa7\x0f(N\xa1\x0b\x05\x19@<T>L9!\xb7\x9e3\xbc\x99\xf0\x8f\xb3\xacZ:\xb3\x1c\xb9\xb7;\xc7\x8a:\xb7\x10\xbd\x07"\xad\xc5\x84'
 
= 206109322179011817882783419945552366363
= 17175776848250984823565284995462697197
= 103054661089505908941391709972776183181
 
# solve for mod p
p_g = complex_pow(g, q * r, n)
p_enc = complex_pow(public_key, q * r, n)
 
# norm : a^2 + b^2
c_1 = norm(p_g) % (p * p)
c_2 = norm(p_enc) % (p * p)
 
c_1 = (c_1 - 1// p
c_2 = (c_2 - 1// p
 
val_p = (c_2 * inverse(c_1, p)) % p 
 
# solve for mod r
r_g = complex_pow(g, p * q, n)
r_enc = complex_pow(public_key, p * q, n)
print(norm(r_g) % p, norm(r_enc) % p)
 
'''
p = 206109322179011817882783419945552366363
g = GF(p)(176015758946526802279559144270141551487) # r_g
enc = GF(p)(28369875517706698292997652748535456248) # r_enc
print(g.multiplicative_order()) # this equals r
print(enc.log(g)) # 26176203815975575469683683780455489251
'''
 
val_r = 26176203815975575469683683780455489251
val_tot, pr = CRT(val_p, p, val_r, r)
 
for i in range(0100):
    private_key = long_to_bytes(val_tot + i * pr)
    flag = AES.new(private_key, AES.MODE_ECB).decrypt(encrypted_flag)
    if b"TetCTF" in flag:
        print(flag)
        break
 
# TetCTF{h0m0m0rph1sm_1s_0ur_fr13nd-mobi:*100*231199111007#}
cs


'CTF' 카테고리의 다른 글

CryptoHack All Solve, Again  (10) 2021.01.31
0x41 CTF SPN Writeup  (0) 2021.01.31
PBCTF 2020 Crypto Writeups  (1) 2020.12.07
N1CTF 2020 Crypto Write-Ups  (0) 2020.10.18
SECCON 2020 OnlineCTF Crypto Write-Ups  (0) 2020.10.11

관련 코드는 github에서 찾아볼 수 있다.


이 글에서는 다음 4가지 문제를 다룬다. 모두 "유클리드 알고리즘에서 등장한 문제 축소 방식"을 이용한다.

  • Lattice Point Counting. $ax+by \le c$, $x, y \ge 1$을 만족하는 격자점 $(x, y)$의 개수를 세는 법.
  • Linear Inequality with Modulo. $L \le Ax \pmod{M} \le R$을 만족하는 최소의 음이 아닌 정수 $x$를 구하는 법.
  • Modulo Maximization. $1 \le x \le C$이며 $Ax \pmod{M}$이 최댓값이 되도록 하는 $x$를 구하는 법.
  • Minimal Fraction. $A/B < x/y < C/D$이며, $y$가 최소인 분수 $x/y$를 찾는 법. 단 $A, C, D \ge 0$, $B \ge 1$.
  • 그리고 사실 이 내용에서 중요한 사례라고 볼 수 있는, continued fraction, 연분수를 매우 간단하게 소개한다. 연분수에 대한 내용은 정수론 책을 찾아보거나, 링크에 걸린 글을 읽는 것을 추천. 필자가 연분수에 대한 내용을 직접 소개하기에는 지식이나 엄밀한 이해/직관이 부족한 것 같다. 


Lattice Point Counting

  • $ax + by \le c$, $x, y \ge 1$을 만족하는 격자점 $(x, y)$의 개수를 세자. 이를 $(a, b, c)$에 대한 문제라고 부르겠다.
  • 일반성을 잃지 않고, $a \ge b$를 가정한다. $a = qb+r$인 $q \ge 1$과 $0 \le r < b$가 존재한다.
  • $ax + by \le c$는 $(qb+r)x + by \le c$, 즉 $b(qx+y) + rx \le c$와 같다.
  • 그러므로, $bu+rv \le c$이며 $u \ge qv+1$, $v \ge 1$를 만족하는 격자점 $(u, v)$의 개수를 세면 된다. ($x=v$, $y=u-qv$에 대응된다)
  • $u$를 고정했을 때 $v \le (c-bu)/r$과 $v \le (u-1)/q$라는 두 부등식이 있다. 둘 중 어느 것이 더 "강한 부등식"인지를 기준으로 경우를 나눈다.
  • 계산을 열심히 하면, 그 기준이 $u = \lfloor cq/a \rfloor$임을 알 수 있다. 이제 경우를 나눠보자.
  • $u > \lfloor cq/a \rfloor$인 경우, $v \le (c-bu)/r$만 만족하면 된다. 그러니 $bu+rv \le c$만 생각하면 된다.
  • $b(u-\lfloor cq/a \rfloor) + rv \le c - b \lfloor cq/a \rfloor$이어야 하며, 원하는 조건은 $u-\lfloor cq/a \rfloor , v \ge 1$이다.
  • 이 경우는, 문제를 $(b, r, c - b \lfloor cq/a \rfloor)$에 대한 문제로 축소시켰다.
  • $u \le \lfloor cq/a \rfloor$인 경우, $v \le (u-1)/q$만 만족하면 된다. 이 식은 $qv + 1 \le u \le \lfloor cq/a \rfloor$과 같다.
  • $v=1$이라면, 가능한 $u$의 개수는 $\lfloor cq/a \rfloor - q$개다. $v$가 $1$ 증가할 때마다, $u$로 가능한 개수가 $q$개 감소한다.
  • 결국 이 경우에는, 등차수열의 합으로 가능한 경우의 수를 표현할 수 있다. $\mathcal{O}(1)$ 선으로 정리할 수 있다.

그러니 문제를 유클리드 호제법과 동일한 시간복잡도로 해결할 수 있다. 

$\sum_{k=1}^n \lfloor (a + bk) / c \rfloor$ 형태의 값 역시 같은 알고리즘으로 계산할 수 있다.

  • 먼저 $a = qc + r$이면, $\lfloor (a+bk) / c \rfloor = q + \lfloor (r+bk)/c \rfloor$이다. 
  • 그러니, $qn$을 값에 더한다음 $a$를 $r$로 바꿀 수 있다. 이제부터 $0 \le a < c$인 경우만 봐도 무방.
  • 위 식은 $1 \le x \le n$, $1 \le y$이며 $a+bx \ge cy$인 격자점 $(x, y)$의 개수를 센다.
  • 이 식을 변형하여, $cy+b(n+1-x) \le (n+1)b+a$ 형태를 얻는다. 
  • $t = n+1-x$라 하면, 목표는 $1 \le t \le n$과 $cy+bt \le (n+1)b + a$다.
  • 애초에 $t \ge n+1$일 수 없는게, 그러면 $cy+bt \ge c + b(n+1) > a + b(n+1)$이 성립한다.
  • 그러니, 단순히 $cy + bt \le (n+1)b + a$, $y, t \ge 1$인 $(y, t)$의 개수를 세면 끝이다. 즉, $(c, b, (n+1)b+a)$ 문제를 풀자.


Linear Inequality with Modulo

  • $L \le Ax \pmod{M} \le R$을 만족하는 최소의 음이 아닌 정수 $x$를 구하는 법을 찾자. 이를 $(A, M, L, R)$ 문제라 하자.
  • 먼저 해가 존재하는지 찾자. $Ax \pmod{M}$은 $\text{gcd}(A, M)$의 배수 전체를 표현할 수 있다. (2단원 내용)
  • 그러니 $L$과 $R$ 사이에 $\text{gcd}(A, M)$의 배수가 존재하는지만 확인하면 된다. 
  • 만약 $L > R$이라면, 구간이 "한 바퀴를 돌았다"고 생각하여, $0$을 문제의 답으로 본다.
  • 즉, $[L, R]$이라는 구간을 수평선의 구간으로 보지 말고, 둘레가 $M$인 원의 한 구간으로 보는 것이다.
  • 이제 본격적으로 문제를 해결하자. 먼저 $0$이 조건을 만족하는지 확인하자. $L > R$이나 $L=0, R=0$인지 보면 된다. 
  • 만약 $2A > M$이라면, $L \le Ax \pmod{M} \le R$과 $M-R \le (M-A)x \pmod{M} \le M-L$이 동치임을 이용하자.
  • 이를 사용하면, 문제를 $(M-A, M, M-R, M-L)$ 문제로 옮길 수 있다. 이제 $2(M-A) < M$이다. 이제부터 $2A \le M$을 가정한다.
  • 만약 $L \le Ax \le R$인 (단, $L, R$ 모두 $[0, M)$에 속한다. $Ax$는 $M$으로 나눈 나머지로 보지 않는다.) $x$가 있다면, 그 중 가장 작은 것을 반환하자.
  • 그렇지 않다는 것은, 일단 $At < L \le R < A(t+1)$인 음이 아닌 정수 $t$가 있음을 의미하고, $R-L+1 < A$임을 얻는다. 이제 진짜 시작. 
  • $L \le Ax \pmod{M} \le R$인 최소 $x$를 찾는 건, $L + My \le Ax \le R + My$인 정수 $x$가 있도록 하는 최소 음이 아닌 정수 $y$를 찾는 것이다.
  • 이는 다시 $L \le Ax- My \le R$과 같은데, 이걸 $\pmod{A}$로 보면 $L \pmod{A} \le (-My) \pmod{A} \le R \pmod{A}$이다.
  • 그러니 $y$의 값을 구하기 위해서는 $(- M \pmod{A}, A, L \pmod{A}, R \pmod{A})$ 문제를 풀어야한다.
  • $y$의 값을 구하면, $x$의 값을 구하기 위해서는 단순히 $L + My \le Ax$인 최소의 $x$를 찾기만 하면 된다.

우리는 modulus에 해당하는 값을 $M$에서 $A$로 줄였으며, $2A \le M$이니 이는 최소 반 줄인 것이다. 그러니 로그 시간.

  • 최소의 음이 아닌 정수 $x$가 아니라, $T$ 이상의 정수 $x$로 조건을 바꿔도 해결할 수 있다.
  • 이 경우, $x = y + T$라고 쓰고 조건을 $(L-AT) \pmod{M} \le Ay \pmod{M} \le (R-AT) \pmod{M}$이라 바꾸자.
  • 이제 위 조건을 만족하는 최소의 음이 아닌 정수 $y$를 구한 뒤, $x = y+T$를 대입하면 끝이다.
  • 이를 반복하면, $L \le Ax \pmod{M} \le R$을 만족하는 자연수 $x$를 순서대로 나열할 수도 있다.


Modulo Maximization/Minimization

  • $1 \le x \le C$를 만족시키면서 $Ax \pmod{M}$이 최댓값이 되도록 하는 $x$를 구하자. 이를 MAX-$(A, M, C)$ 문제라 하자.
  • $1 \le x \le C$를 만족시키면서 $Ax \pmod{M}$이 최솟값이 되도록 하는 $x$를 구하자. 이를 MIN-$(A, M, C)$ 문제라 하자.
  • MAX 문제나 MIN 문제나 사실 같은 문제다. $A$를 $-A$로 바꿔서 생각하면 두 문제가 같은 문제가 되기 때문이다.
  • 이 문제는 "Linear Inequality with Modulo" 문제에 이분탐색을 추가해서 풀 수 있다.
  • adamant의 블로그 글이 시각화도 좋고, 설명도 좋고, 추가할 점도 특별히 없는 것 같다. 링크를 거는 것으로 대체한다.
  • 그런 줄 알았는데, 실제로 코드를 비교해보니 약간의 문제가 있는 것으로 보인다. 일단은 링크의 아이디어만 알아두자.


Minimal Fraction

  • $A/B < x/y < C/D$를 만족시키면서, $y$가 최소인 분수 $x/y$를 찾는 경우. 이를 $(A, B, C, D)$ 문제라 하자.
  • 위에서도 언급했지만, $A, C, D \ge 0$, $B \ge 1$을 가정한다. 우리의 목표 $x, y$ 역시 음이 아닌 정수여야 한다.
  • 사실 $y$를 최소화하는 것이나 $x$를 최소화하는 것이나 큰 차이가 없다. $y$가 결정되었다면, $Ay/B < x < Cy/B$인 최소의 $x$를 잡으면 된다. 
  • $y$가 커지면, 자연스럽게 가능한 $x$의 범위도 "오른쪽으로" 이동하게 될 것이다.
  • $D=0$인 경우, $C/D$를 $\infty$로 간주하도록 하겠다. 이 경우까지 처리해야 문제가 풀린다.
  • 이 문제는 아래에 언급할 Stern-Brocot Tree를 이용한 풀이가 있다. 링크가 걸린 myungwoo님의 글을 참고하라.
  • $A/B < C/D$가 성립하지 않는 경우, 답은 당연히 없고, 이 결과를 반환한다.
  • $A=0$인 경우. 이때는 $0 < x/y < C/D$인 것만 생각하면 된다. $x=1$로 고정하고, $1/y < C/D$인 최소의 $y$를 잡으면 ok.
  • $A/B$와 $C/D$ 사이에 (양 끝 제외) 정수가 존재하는 경우. $y=1$을 정하고 $x$를 존재하는 정수 중 가장 작은 것으로 잡으면 ok.
  • $C/D$가 정수인 경우. 이때, $C/D=n$이라 하자. 이 경우, $A/B < n-1/y$인 최소 $y$를 잡고, $Ay/B < x < ny$인 최소 $x$를 잡으면 ok.
  • 이제 $n \le A/B < C/D < n+1$인 자연수 $n$이 존재하는 경우만 남았다. 이 단계가 핵심.
  • 목표는 $A/B < x/y < C/D$인 최소 $x, y$를 찾는 것이다. $n$을 빼준 다음, 전부 역수를 취해주자.
  • 이제 목표는 $D/(C-nD) < y/(x-ny) < B/(A-nB)$인 최소의 $x,  y$를 찾는 것이다.
  • 애초에 분모를 최소화하는 것이나 분자를 최소화하는 것이나 같은 문제이므로, $(D, C-nD, B, A-nB)$ 문제를 풀면 된다.
  • 여기서 $n$이 $A, C$를 $B, D$로 나눈 몫임을 생각하면, 이 과정이 유클리드 호제법과 같음을 알 수 있다. 끝. 

관련 문제로 Google Code Jam 2019 Round 2 C번 문제가 있다. 아래 myungwoo님의 글도 그 당시 나왔다.


Continued Fraction

  • $[a_0 ; a_1, \cdots , a_k] = a_0 + (1 / (a_1 + (1 / (a_2 + \cdots ) ) ) ) $라고 정의하자. 이를 연분수라고 부른다.
  • 모든 유리수는 유한한 연분수로 나타낼 수 있다. $a=qb+r$이라 하면, $a/b = q + r/b = q + 1/(b/r)$이다.
  • 이는 $(a, b)$에 대한 문제를 $(b, r)$에 대한 문제로 줄인 것과 같다. 유클리드 호제법 그 자체...
  • 연분수는 정말 다양한 성질들이 있다. 이는 학부 정수론 책을 참고하는 것을 추천.
  • 또한, 위에서도 언급한 adamant의 Codeforces 블로그를 참고하는 것을 추천한다. 좋은 자료가 많다.
  • 연분수와 밀접한 연관이 있는 개념으로 Stern-Brocot Tree가 있다. 이에 대한 글은 myungwoo님의 이 글을 참고하라. 
  • 위 링크에서도 나와있지만, Stern-Brocot Tree로 유리수에서 효율적인 이분탐색을 하는 등 매우 흥미로운 알고리즘들이 많다.


관련 코드는 github에서 찾아볼 수 있다.


이 글의 순서는 다음과 같다.

  • 원시근의 정의 및 관련 사실들. 핵심 중의 핵심을 제외하고, 증명 대부분은 생략할 것이다.
  • 원시근을 찾는 알고리즘과 위수를 계산하는 알고리즘.
  • 이산로그 문제와 Baby Step Giant Step. 그리고 그 확장인 Pohlig-Hellman 알고리즘 (optional!)
  • 이산제곱근 문제와 그 해결법. 특수한 경우인 modular square root의 경우. 그리고 Adleman-Manders-Miller (매우 optional!) 
  • 특히, Adleman-Manders-Miller는 지금까지 나온 경우를 딱 한 번 봤다. (GP of Baltic Sea) 이는 논문 링크 + 구현체 링크로 대체.
  • 그리고 마지막으로, FFT와 NTT 사이의 연결고리를 설명한다. 앞서 mobius inversion = 포함-배제를 설명한 것과 같이, FFT = NTT를 설명한다.


원시근이란 무엇인가? 위수란 무엇인가?

  • 자연수 $n \ge 2$이 있을 때, $n$과 서로소이며 $n$ 미만인 자연수의 개수는 $\phi(n)$이다.
  • 만약 $g, g^2, \cdots , g^{\phi(n)}$을 각각 $n$으로 나눈 나머지가 이 $\phi(n)$개의 자연수 전부가 된다면, $g$를 $n$의 원시근이라 한다.
  • 다른 말로 하면, $t$가 $n$과 서로소인 자연수라면 적당한 $1 \le k \le \phi(n)$이 있어 $g^k \equiv t \pmod{n}$이란 것이다.
  • 사실 1. 원시근을 갖는 자연수는 $2, 4$와 홀수 소수 $p$와 자연수 $k \ge 1$에 대하여 $p^k$, $2p^k$다.
  • 사실 2. $n$이 원시근을 갖는 자연수라면, 원시근은 정확히 $\phi(\phi(n))$개 있다. 다른 말로 하면, "원시근은 꽤 많다".
  • 예를 들어, $n=p$가 소수라면, $\phi(\phi(p))=\phi(p-1)$이다. 특히, $\phi(n) > n/(e^{\gamma} \log \log n + 3/\log \log n)$임이 알려져 있다.
  • 자연수 $g, n$이 있을 때, (단, $\text{gcd}(g, n)=1$) $g^k \equiv 1 \pmod{n}$을 만족하는 최소 자연수 $k$를 $g$의 위수라고 한다.
  • 사실 3. $g$가 $n$의 원시근인 것은, $g$의 위수가 $\phi(n)$인 것과 동치이다. (증명 힌트. $g^i \equiv g^j \pmod{n}$이면 $g^{j-i} \equiv 1 \pmod{n}$) 
  • 사실 4. ("학부 대수학의 반의 avatar") $\text{gcd}(g, n)=1$이고 $g$의 위수가 $k$라고 하자. $g^m \equiv 1 \pmod{n}$이 성립할 필요충분조건은 $k|m$이다. 
  • 사실 4 증명. $k|m$이면 $g^m \equiv (g^k)^{m/k} \equiv 1^{m/k} \equiv 1 \pmod{n}$이므로 ok. 만약 $g^m \equiv 1 \pmod{n}$인데 $k|m$이 아니라고 하자. 그러면 $m=kt+r$인 $0 < r < k$가 있고, $g^r \equiv g^{m-kt} \equiv 1 \pmod{n}$이다. 그런데 $0 < r < k$이므로 이는 위수의 "최소성"에 모순이다. 즉, 애초에 $k$가 $g^k \equiv 1 \pmod{n}$을 만족하는 "최소의" 자연수였는데, $g^r \equiv 1 \pmod{n}$, $r<k$이므로 모순. 좋은 아이디어.
  • 사실 5. $g$의 위수가 $t$라면, $g^d$의 위수는 $t/\text{gcd}(d, t)$가 된다. 사실 4에서 바로 증명할 수 있다.
  • 사실 1, 2의 증명은 학부 정수론에서 가장 중요한 부분 중 하나이며, 후에 (사실 1은) 현대대수학에서도 다시 증명하게 된다.


왜 원시근이 강력한가?

  • $g$가 원시근이라 하자. $g^k \equiv r \pmod{n}$이라면, $k$를 $r$의 "로그 값"이라고 생각할 수 있겠다.
  • 로그가 강력한 이유가 뭐였더라? 거듭제곱을 곱셈으로, 곱셈을 덧셈으로 바꿔주는 것이었다.
  • 우리는 거듭제곱은 잘 모르지만, 앞서 2, 3 단원을 통해서 일차합동식 $ax \equiv b \pmod{n}$에 대해서는 안다.
  • 그러니 원시근을 사용해서 거듭제곱에 관한 문제를 일차식으로 바꿀 수 있을 것이다. 이 접근을 "이산제곱근"에서 사용한다.


이제부터 원시근을 찾는 방법을 설명한다. $n$이 원시근을 갖는 자연수, 즉 $2, 4, p^k, 2p^k$ 형태라고 가정하자.

  • 원시근은 충분히 많으므로, 랜덤하게 $g$를 잡은 뒤 $g$가 원시근인지 판별하는 과정을 원시근을 찾을 때까지 반복하면 된다.
  • 원시근이 많다는 것은, 한 번 랜덤한 시도를 했을 때 성공적으로 원시근을 찾을 확률이 나쁘지 않다는 것이다.
  • $g$가 원시근인지 판별하려면, 먼저 $\text{gcd}(g, n)=1$인지 확인한다. 그 후, $g$의 위수가 $\phi(n)$인지 확인하면 된다.

그러므로, 원시근을 찾는 방법을 알기 위해서는 위수를 계산하는 알고리즘을 찾으면 된다. 여기서 사실 4가 활용된다.

  • $\phi(n)$의 소인수분해가 이미 계산되었다고 가정한다. $\phi(n) = p_1^{e_1} \cdots p_k^{e_k}$라고 가정한다.
  • $n=p^k$, $n=2p^k$인 경우 $\phi(n) = p^{k-1}(p-1)$임에 유의하라. 
  • 이제 $g$의 위수를 찾는 것이 목표다. $\text{gcd}(g, n)=1$임을 가정한다. 오일러의 정리에 의하여, $g^{\phi(n)} \equiv 1 \pmod{n}$.
  • 이는 사실 4에 의하여, 위수의 값이 $\phi(n)$의 약수임을 의미한다. 
  • 이제 $g^{\phi(n) / p_1}$의 값을 보자. 만약 이 값이 $1 \pmod{n}$이라면, 사실 4에서 위수가 $\phi(n)/p_1$의 약수임을 얻는다.
  • 만약 $g^{\phi(n)/p_1} \equiv 1 \pmod{n}$이었다면, $g^{\phi(n)/p_1^2}$을 보자. (단, $p_1^2 | \phi(n)$) 이 값 역시 $1 \pmod{n}$이라면, 사실 4에서 위수가 $\phi(n) / p_1^2$의 약수.
  • $g^{\phi(n)/p_1^2}$은 $1 \pmod{n}$이 아니었다고 하면, 위수가 $\phi(n)/p_1$의 약수였으나 $\phi(n)/p_1^2$의 약수는 아님을 의미한다.
  • 이는 다르게 말하면 위수가 $p_1$을 몇 개 가지고 있는지 파악했다는 이야기다. 이제 $p_2$로 넘어가서 같은 작업을 반복, $p_k$까지 하자.
  • 즉, $\phi(n)$에서 시작하여, $g^{???} \equiv 1 \pmod{n}$이 성립할 때까지 계속 소인수들을 제거해주면 된다.
  • 시간복잡도를 생각해보면, 대강 $(e_1+e_2 + \cdots + e_k) = \mathcal{O}(\log n)$번  $g$의 거듭제곱을 계산하므로, $\mathcal{O}(\log^2 n)$이다.
  • 물론, $\phi(n)$의 소인수분해가 필요했음에 주의하라. 필요한 경우, 6단원에서 배운 Pollard-Rho 소인수분해를 활용하자. 

 

이산로그 문제와 Baby Step Giant Step

  • 원시근 $g$와 $\text{gcd}(h, n)=1$을 만족하는 자연수 $h$가 주어졌을 때, $g^x \equiv h \pmod{n}$을 만족하는 $x$를 찾는 문제를 이산로그 문제라 한다.
  • 즉, 앞서 언급했던 원시근의 강력함인 "로그"를 실제로 계산하는 문제이다. 이산적인 집합에서 계산하고 있으니 이산로그 문제다.
  • 이 문제를 CP/PS에서 보이는 "meet-in-the-middle"을 이용하여 $\mathcal{O}(\sqrt{n} \log n)$ 시간에 해결하는 알고리즘이 있다. 
  • $B = \lceil \sqrt{\phi(n)} \rceil$이라고 하자. 그러면 $0$과 $\phi(n)$ 사이의 모든 값은 적당한 $0 \le u, v < B$에 대하여 $uB+v$라고 쓸 수 있다.
  • 그러니 우리는 $g^{uB+v} \equiv h \pmod{n}$을 만족하는 $0 \le u, v < B$를 찾으면 된다. 
  • 그런데 이 식을 변형하면 $(g^B)^{u} \equiv h \cdot g^{-v} \pmod{n}$을 얻는다. 이제 meet-in-the-middle 각이 보인다.
  • $h \cdot g^{-0}$부터 $h \cdot g^{-(B-1)}$를 C++의 std::map과 같은 자료구조에 넣어놓자. 이때, key 값을 $h \cdot g^{-v}$로 하고, value 값을 $v$로 하도록 한다. 
  • 이제 $u$에 대하여 순회하면서, $(g^B)^{u} \pmod{n}$이 해당 자료구조의 key 값에 있는지 확인하자. 있다면, value 값에서 $v$를 얻는다.
  • 이 방식으로 $u, v$를 찾을 수 있으며, 시간복잡도는 $\mathcal{O}(\sqrt{n} \log n)$이다. 로그는 C++의 std::map 때문에 붙는다.
  • 이 알고리즘은 이렇게 $\pmod{n}$에 대한 이산로그를 구하는 것에서만 쓰일 수 있는 것이 아니라, finite abelian group 구조를 갖는 모든 문제에서 적용될 수 있다. 애초에 이산로그 문제 자체가 finite abelian group 구조에서 정의되는 문제이며, 지금 우리가 다루고 있는 문제는 그 중 일부인 경우이다. 
  • 말은 어렵게 했으나, 실제로 문제를 풀면 다른 생각이 들 수도 있다. BOJ의 xorshift32 문제를 참고하라.
  • Note. $g$가 원시근이 아니어도, $g$의 위수를 $k$라고 하면 같은 방법으로 $\mathcal{O}(\sqrt{k} \log k)$ 알고리즘을 생각할 수 있다. 위에서 $\phi(n)$을 $k$로 바꾸면 된다.



Pohlig-Hellman 알고리즘 (optional)

  • 위 Baby-Step-Giant-Step 알고리즘을 강화한 것이 Pohlig-Hellman 알고리즘이다. $\phi(n) = p_1^{e_1} \cdots p_k^{e_k}$라고 하자. 
  • $g^x \equiv h \pmod{n}$인 $x$를 찾는다는 것은, 사실 $x \pmod{\phi(n)}$을 찾는다는 것이다.
  • 그러니, 중국인의 나머지 정리를 생각하면 $x \pmod{p_i^{e_i}}$들을 계산한 다음 합치면 충분하다.

여기서 $\mathcal{O}(\sqrt{p_i} \log p_i)$ 시간에 $x \pmod{p_i}$를 찾는 방법을 제시한다. 

$g^x \equiv h \pmod{n}$이면, $(g^{\phi(n)/p_i})^x \equiv h^{\phi(n)/p_i} \pmod{n}$ 역시 성립한다. 여기서 $g^{\phi(n)/p_i}$의 위수가 $p_i$임에 주목하자.

이 새로운 이산로그 문제를 풀면 $x \pmod{p_i}$를 얻고, 이를 위해서 소모되는 시간은 $\mathcal{O}(\sqrt{p_i} \log p_i)$가 된다.

이제 이를 확장하여, $x \pmod{p_i^2}$을 찾는 방법을 고안하자. 물론, $p_i^2 | \phi(n)$이라고 가정한다.

마찬가지로, $(g^{\phi(n)/p_i^2})^x \equiv h^{\phi(n)/p_i^2} \pmod{n}$을 생각한다. $g^{\phi(n)/p_i^2}$의 위수는 $p_i^2$이다.

그러니 이 이산로그 문제를 풀면 $x \pmod{p_i^2}$을 얻는데, 중요한 점은 우리가 이미 $x \pmod{p_i}$를 안다는 것이다.

이미 계산한 값인 $x \pmod{p_i}$를 $r$이라고 하자. 이제 $x \pmod{p_i^2} = p_ik + r$이라고 쓸 수 있고, 목표는 $k$가 된다. 그러면 $$ (g^{\phi(n)/p_i})^k \equiv h^{\phi(n)/p_i^2} \cdot (g^{\phi(n)/p_i^2})^{-r} \pmod{n}$$을 얻는다. 이제 $k$를 찾기 위해 이산로그 문제를 풀면 $\mathcal{O}(\sqrt{p_i} \log p_i)$ 시간에 $x \pmod{p_i^2}$을 얻는다. 

이를 $p_i^{e_i}$까지 반복하고, 이를 각 prime power에 대해 반복하면 끝. 시간복잡도는 $\mathcal{O}(\sum e_i \sqrt{p_i} \log p_i )$다.

관련 문제로는 위에 언급한 xorshift32를 강화한 xorshift64가 있다. 


그리고 필자도 굉장히 최근에 배운 트릭을 하나 소개한다. 매우 optional인 건 맞지만, 흥미롭고 Pohlig-Hellman 보다 쉽다.

  • $g^x \equiv h \pmod{p^2}$를 푼다고 가정하자. 단 $g$는 $p^2$에 대한 원시근이다. 
  • 먼저 $g^{a_1} \equiv h \pmod{p}$를 푼다. 그러면 여기서 얻는 결과는 $x \equiv a_1 \pmod{p-1}$이 될 것이다.
  • 이제 $g^{a_2} \equiv h \pmod{p^2}$을 푼다. 이제 $a_2 \equiv a_1 \pmod{p-1}$을 알고 있으므로, $a_2 \pmod{p}$를 알면 된다.
  • 이를 위해서, $g^{(p-1)a_2} \equiv h^{p-1} \pmod{p^2}$을 푼다. 그런데, 놀라운 점은 $g^{p-1}$, $h^{p-1}$ 모두 $1 \pmod{p}$라는 것이다.
  • 그래서 사실 $(1+pu)^{a_2} \equiv (1+pv) \pmod{p^2}$ 형태의 식을 풀게 되며, 특히 $g$가 원시근이면 $u$는 $p$의 배수가 아니다.
  • 또한, 이항정리에서 $(1+pu)^{a_2} \equiv 1+pua_2 \pmod{p^2}$이다. 그러니 결국 $ua_2 \equiv v \pmod{p}$를 풀면 ok.
  • 그러니 Baby-Step-Giant-Step을 두 번 돌릴 필요가 없으며, 한 번만 돌려도 된다. 이는 확장이 가능한 아이디어다. 이 링크 참고.


이산제곱근 문제

  • 이번에는 $x, b$가 주어졌을 때, (단, $\text{gcd}(b, n)=1$) $a^x \equiv b \pmod{n}$을 만족하는 $a$를 구하는 것이 목표다.
  • 중국인의 나머지 정리의 아이디어를 가져오자. $n = p_1^{e_1} \cdots p_k^{e_k}$라 하면, $a \pmod{p_i^{e_i}}$를 각각 구한 뒤 합치는 전략을 세울 수 있다.
  • 그러니 $n$이 prime power인 경우에서만 해결하면 충분하다. 
  • 여기서는 $n$이 홀수 소수의 거듭제곱인 경우만 다루겠다. 이제 원시근 $g$의 존재를 가정해도 된다.

$a = g^u$, $b = g^v$인 $u, v$를 이산로그 알고리즘으로 찾자. 목표는 $g^{ux} \equiv g^v \pmod{n}$인 $x$를 구하는 것이다.

이는 $g^{ux-v} \equiv 1 \pmod{n}$과 같으니, 사실 4에 의하여 $ux \equiv v \pmod{\phi(n)}$을 푸는 것과 같다.

이는 일차합동식이므로, 2단원에서 배운 내용을 통해 해결할 수 있다. $u$를 알면, 이를 가지고 $a \pmod{n}$ 역시 구할 수 있다.


특히, $\text{gcd}(x, \phi(n))=1$이라면, $tx \equiv 1 \pmod{\phi(n)}$인 $t$를 찾을 수 있다.

그러면 $b^t \equiv a^{tx} \equiv a \pmod{n}$이 성립함을 알 수 있고, (오일러 정리) $a \pmod{n}$을 쉽게 찾을 수 있다.


위 예시에서 $2$의 거듭제곱이 여러모로 까다롭다는 것을 알 수 있다. 이렇게 $2$의 거듭제곱을 가지고 귀찮게하는 문제는 적다.

  • 하지만, 이 사실이 도움이 될 수 있다 : 모든 $x \pmod{2^n}$은 (단, $x$ 홀수) $5^k$ 또는 $-5^k$로 나타낼 수 있다.
  • 즉, $x$가 $1 \pmod{4}$인 경우, 전부 $5$의 거듭제곱으로 표현이 되며, $3 \pmod{4}$인 경우, $-x$가 $5$의 거듭제곱으로 표현이 된다.
  • 그러니 $5$를 $2^n$의 일종의 원시근과 비슷한 무언가로 생각할 수 있겠다. 이를 통해 이산로그 등 문제를 접근할 수 있다.


modular square root

  • 이 파트에서는 특수한 경우인, $x^2 \equiv a \pmod{n}$을 다룬다. 편의상 $\text{gcd}(a, n) = 1$을 가정한다.
  • 역시 중국인의 나머지 정리 아이디어를 사용하여, $x^2 \equiv a \pmod{p^k}$만 해결하면 된다.

먼저 $p$가 홀수 소수인 경우부터 보자. 특히, $k=1$인 경우, 즉 $x^2 \equiv a \pmod{p}$인 경우를 보자.

  • 이차잉여. 위 식이 해를 가질 조건은 $a^{(p-1)/2} \equiv 1 \pmod{p}$이다. 
  • 증명. 스케치만 하자면, 해를 가질 조건이 $a$가 $g^{even}$인 것과 동치이며 이건 다시 $a^{(p-1)/2} \equiv 1 \pmod{p}$와 동치이다.

해를 갖는 경우, $\mathcal{O}(\log^2 p)$ 시간에 해를 찾는 알고리즘이 있다. Tonelli-Shanks 알고리즘이며, 설명은 이 블로그에 있다.

특히, $x$가 해면 $-x$ 역시 해임에 주목하자. 해는 정확히 2개 있을 것이다. (이유가 궁금하다면, 이 정리 참고) 이제 문제는 $k \ge 2$인 경우이다. 


이제부터 각 $x^2 \equiv a \pmod{p^k}$의 해를 $x^2 \equiv a \pmod{p^{k+1}}$의 해로 "lift"하는 방법을 생각해보자.

이미 $r^2 \equiv a \pmod{p^k}$인 $r$의 값을 찾았다고 생각하자. 해는 정확히 $r, -r$ 두 개가 있을 것이다. 이제 $x^2 \equiv a \pmod{p^{k+1}}$을 풀자.

$x^2 \equiv a \pmod{p^k}$ 역시 성립하니, $x \equiv r \pmod{p^k}$ 또는 $x \equiv -r \pmod{p^k}$가 성립해야 한다.

  • 이제 $x \equiv r \pmod{p^k}$인 경우를 가정하자. 그러면 $x \pmod{p^{k+1}}$은 $r + t p^k$ 형태로 쓸 수 있다. 단, $0 \le t < p$. 
  • 이제 $x^2 \pmod{p^{k+1}}$을 생각하면, $r^2 + 2rt p^k + t^2 p^{2k} \equiv r^2 + 2rt p^k \pmod{p^{k+1}}$과 같다.
  • 그러니, 문제는 결국 $r^2 + 2rt p^k \equiv a \pmod{p^{k+1}}$이고, 이를 변형하면 $2rtp^k \equiv (a-r^2) \pmod{p^{k+1}}$이다. 
  • 여기서 $a-r^2 \pmod{p^{k+1}}$이 $p^k$의 배수임에 주목하자. 
  • 그러니 원하는 것은 $(2rt - (a-r^2)/p^k) \cdot p^k$가 $p^{k+1}$의 배수가 되는 것, 즉 $2rt \equiv (a-r^2)/p^k \pmod{p}$이다. 
  • 2와 $r$이 전부 $p$와 서로소이므로, 이 합동식을 풀어 $t \pmod{p}$를 얻는다. 이제 $x = r + tp^k$가 해이다.
  • 시작을 $x \equiv -r \pmod{p^k}$로 얻었다면, $-r-tp^k$를 얻었을 것이다. 그러니 여전히 해는 $x, -x$ 정확히 두 개다.

이러한 과정을 잘 나타내는 정리가 Hensel's Lemma다. 이제 $p=2$인 경우를 생각해보자.


$k=1, 2$인 경우, 즉 $\pmod{2}$, $\pmod{4}$에서 식을 푸는 것은 brute force로 가능할 것이다. 이 경우는 생략.

이제부터 $k \ge 3$을 가정한다. 핵심적인 사실은, $x^2 \equiv 1 \pmod{8}$이 모든 홀수 $x$에 대해서 성립한다는 것이다.

  • $x^2 \equiv a \pmod{2^k}$가 해를 가질 필요충분조건은 $a \equiv 1 \pmod{8}$.
  • 이때, 해는 정확히 $4$개 존재한다. ($\pmod{2^k}$에서) $x$가 해라면 $x, -x, x+2^{k-1}, -x+2^{k-1}$이 해가 된다.

역시 "lift"할 각을 재보자. $x^2 \equiv a \pmod{2^3}$의 해는 (애초에 $a \equiv 1 \pmod{8}$이니) $x = 1, 3, 5, 7 \pmod{8}$이다.

해가 4개가 있고, 하나만 알면 나머지 세 개의 해를 쉽게 얻을 수 있다. 그러니 해를 하나만 관리해도 충분할 것이다. $x = 1$로 시작하자.

$x^2 \equiv a \pmod{2^k}$의 해를 하나 찾았다고 하고, 이를 $r$이라 하자. 그러면 $(r + t2^{k-1})^2 \equiv r^2 + rt 2^k \pmod{2^{k+1}}$이다. 

만약 $r^2 \equiv a \pmod{2^{k+1}}$이라면, $r$이 그대로 $x^2 \equiv a \pmod{2^{k+1}}$의 해가 된다.

그렇지 않다면, $r^2 \equiv a + 2^k \pmod{2^{k+1}}$일 것이다. 그러면 $x = r + 2^{k-1}$이 해가 된다. 이제 끝.



Adleman-Manders-Miller (매우 optional)

  • 이산제곱근을 더욱 효율적으로 계산하는 알고리즘으로, Tonelli-Shanks의 확장이다.
  • 이 글, 이 논문, 그리고 KAIST의 ICPC팀(이었던) 더불어민규당의 팀노트를 (discrete kth root) 참고하자.
  • PS/CP 환경에서는 쓰이는 것은 저 오픈컵을 제외하고 본 적이 없다. 개인적으로도 암호학 대회 환경을 제외하면 쓴 적이 없다.

  

FFT와 NTT

  • FFT와 FFT의 원리를 이미 알고 있다고 가정한다. $2^n = N$이 FFT의 "크기"라고 가정하자.
  • FFT에서 사용되는 사실은 $\omega = e^{2\pi i / N}$이라 할 때, $\omega^{N}=1$이고 $1, \omega, \cdots , \omega^{N-1}$이 모두 다르다는 것이다.
  • 또한, inverse FFT 과정에서는 계산 후에 $N$을 나누어야 하므로, 복소수에서 $N$을 나누는 연산이 잘 정의됨을 이용한다.
  • NTT도 다르지 않다. $p-1 = 2^ab$이고, $a$가 충분히 큰 소수 $p$와 $p$에 대한 원시근 $g$를 준비하자. $2^a = N$이라 하자.
  • $ \omega = g^{b}$라 하면, $\omega^{N} = 1$이고, $1, \omega, \cdots , \omega^{N-1}$이 모두 다르다. (모든 연산은 $p$로 나눈 나머지로 본다)
  • 마찬가지로, $N$이 $p$와 서로소이므로, $N$으로 나누는 연산이 잘 정의된다. (modular inverse가 존재하므로)
  • 그러니, FFT와 NTT는 정확히 같은 알고리즘이다. 실제 코드도 별로 다를 것이 없다.


독자들은 정수론은 잘 모르더라도 포함-배제의 원리는 어느 정도 숙지하고 있을 것이다. 

포함-배제의 원리는 애초에 교육과정의 경우의 수 세는 문제에서도 매우 많이 등장하며, "mainstream CP/PS"에서도 꽤 등장하기 때문이다.

이 파트에서 이야기하고자 하는 것은 mobius function/inversion이 포함-배제와 다를 게 없다는 것이다. 

(사실, 이 사실은 제곱 ㄴㄴ 문제를 해결한 사람이라면 이미 느꼈을 것이라고 생각한다. 여기서 더 자세하게 보자.)


이제부터 포함-배제의 원리 이야기로 넘어간다. 

  • 포함-배제의 원리의 가장 익숙한 형태는, 유한집합 $A_1, A_2, \cdots ,A_n$에 대하여 $$\displaystyle |\cup_{i} A_i| = \sum_{i} |A_i| - \sum_{i< j} |A_i \cap A_j| + \sum_{i < j < k} |A_i \cap A_j \cap A_k| - \cdots  + (-1)^{n-1} | A_1 \cap A_2 \cap \cdots \cap A_n|$$이 성립한다는 명제일 것이다. 특히 $n=2, 3$인 경우는 고등학교 수학. 
  • 위 식을 이해하는 방법은 다양할 것이다. 하지만 개인적인 생각으로 가장 쉬운 방법은 "원소별 기여"를 보는 것이다.
  • $x \in \cup_i A_i$를 아무거나 하나 잡자. 이 원소는 $A_1, \cdots , A_n$ 중 정확히 $k$개의 집합에 속한다고 하자. 물론, $1 \le k \le n$.
  • $x$의 존재로 위 식의 좌변은 $1$만큼 커진다. 그러니, $x$가 우변에 "기여하는 값"이 $1$임을 증명하면 충분하다.
  • $\sum_i |A_i|$에 $x$가 기여하는 횟수는 $k$번, $\sum_{i < j} |A_i \cap A_j|$에 기여하는 횟수는 $x$가 $A_i, A_j$ 각각에 속해야 하므로 $\displaystyle \binom{k}{2}$번.
  • 이제 느낌을 알 것 같다. $x$가 우변에 기여하는 총 값은 $\displaystyle \sum_{i=1}^n (-1)^{i-1} \binom{k}{i}$이다. 단, $n < m$면 $\displaystyle \binom{n}{m} = 0$이라 정의한다.
  • 그런데 이항정리에서 $\displaystyle \sum_{i=0}^n (-1)^i \binom{k}{i} = \sum_{i=0}^k (-1)^i \binom{k}{i} = (1-1)^k =  0$이다. 이를 통해서, $x$가 기여하는 값이 $1$임을 얻는다.


이제 본격적으로 작업을 시작하자. 내용은 이 글을 따라간다.

  • 포함-배제의 원리의 다른 형태를 소개한다. 전체집합 $U$가 유한집합이라 하자. $U$의 각 부분집합 $S$에 대하여 값 $f(S) \in \mathbb{R}$을 주자. 
  • 이제 $U$의 각 부분집합 $A$에 대하여 $g(A) = \sum_{S \subset A} f(S)$를 정의하자. 우리의 목표는 $g$의 값을 가지고 $f$의 값을 복원하는 것이다.
  • 이때, $f(A) = \sum_{S \subset A} (-1)^{|A|-|S|} g(S)$가 성립한다. 이를 위와 똑같은 방법으로 증명해보자.
  • 우변은 $\sum_{S \subset A} (-1)^{|A|-|S|} \sum_{T \subset S} f(T) = \sum_{T \subset S \subset A} (-1)^{|A|-|S|} f(T)$다. 이제 $f(T)$가 "기여하는 값"을 생각해보자. 
  • $T = A$인 경우, $\sum_{A \subset S \subset A} (-1)^{|A|-|S|} = (-1)^0 = 1$이다. 그러니 우변에 $f(A)$라는 값이 추가된다.
  • $T \subsetneq A$인 경우, $\sum_{T \subset S \subset A} (-1)^{|A|-|S|}$를 계산해보자. $|A|=m+n$, $|T|=m$이라고 하면, $T \subset S \subset A$를 만족하는 크기 $m+k$의 집합 $S$는 총 $\displaystyle \binom{n}{k}$개 존재할 것이다. 그러므로, $\displaystyle \sum_{T \subset S \subset A} (-1)^{|A|-|S|} = \sum_{k=0}^n \binom{n}{k} \cdot (-1)^{n-k} = (1-1)^n = 0$이다. 여기서 $n \neq 0$이 사용되었다. 이는 $f(T)$가 사실 우변에 없음을 (0이 곱해져서 더해짐을) 의미한다. 그러니 사실 우변을 정리하면 남는 것은 $f(A)$ 뿐이며, 증명 끝.


위 "새로운 형태"는 포함-배제 느낌이 나지만, 뭔가 갑갑함을 느끼는 독자들이 있을 것이다.

  • 우리에게 익숙한 형태로 돌아가서, $A_1, A_2, \cdots , A_n$이란 유한집합이 있다고 가정하자.
  • 이제 다른 형태로 간다. $U = \{1, 2, \cdots , n\}$이라 하자. 각 $U$의 부분집합 $S$에 대해, $f(S)$를 다음과 같이 정의한다.
  • $f(U) = 0$이다. 만약 $S$가 $U$가 아니라면, $f(S)$는 "속하는 집합의 index 모음이 정확히 $U - S$인 원소의 개수"다.
  • 이게 무슨 뜻이냐? $n=3$이라 하자. $f(\{1\})$은, 정확히 $A_2, A_3$에만 속하는 원소의 개수, 즉 $|A_1^C \cap A_2 \cap A_3|$이다.
  • 이제 $g(A)$를 "포함-배제의 새로운 형태"에서 다룬 방식으로 정의하고, $0 = f(U) = \sum_{S \subset U} (-1)^{|U|-|S|} g(S)$라는 식을 써보자.
  • 이 식이 우리에게 익숙한 포함-배제의 형태와 동치임을 알 수 있을 것이다. 좋은 Exercise이며, CP/PS에서도 충분히 활용될 수 있는 연습이다.
  • 힌트를 주자면, $g(A)$는 "속하는 집합의 index 모음이 $U-A$를 포함하는 원소의 개수"다. 앞서 다룬 $n=3$ 예시에서 $g(\{1\}) = |A_2 \cap A_3|$다.
  • 특히, $g(U) = |\cup_i A_i|$가 됨을 확인할 수 있다. 이제 식을 전부 대입하고 약간의 고민을 하면 원하는 결과를 얻을 것이다.


이제 집합 $X$에 대하여 $\mu(X) = (-1)^{|X|}$라 정의하자.

  • 우리의 결론은 $g(A) = \sum_{S \subset A} f(S)$면, $f(A) = \sum_{S \subset A} (-1)^{|A|-|S|} g(S) = \sum_{S \subset A} \mu(A-S) g(S)$이 성립한다는 것이다.
  • $\mu$를 정의한 상태에서, "새로운 형태"에 대한 증명을 다시 보자. 증명의 핵심은 $$\sum_{T \subset S \subset A} (-1)^{|A|-|S|} = 0 $$이 $T \subsetneq A$일 때 성립함을 보인 것이다. $(-1)^{|A|-|S|} = \mu(A-S)$이고, $T \subset S \subset A$는 $A-S \subset A-T$와 동치이므로, 이 결과는 사실 $$ \sum_{X \subset (A-T)} \mu(X) = 0$$과 같다. 특히, $T \subsetneq A$이므로 $A-T$는 nonempty다. 결과적으로, 증명의 핵심은 모든 nonempty가 아닌 집합 $U$에 대하여 $$\sum_{X \subset U} \mu(X) = 0$$인 것이다. 한편, $U$가 공집합일 때는 좌변의 값이 $1$임 역시 자명하다.


집합이 아니라, 중복 원소가 허용되는 multiset이라면 어떻게 될까? 이때는 $\mu$를 다음과 같이 정의하자.

  • $X$가 중복 원소를 갖는 multiset이라면, $\mu(X) = 0$.
  • $X$가 중복 원소를 가지지 않는다면, 기존의 경우와 같이 $\mu(X) = (-1)^{|X|}$.

이제 여전히 $g(A) = \sum_{S \subset A} f(S)$이면, $f(A) = \sum_{S \subset A} \mu(A-S) g(S)$임을 보이자.

  • 증명의 핵심만 보이면 충분하다. 즉, $\sum_{X \subset U} \mu(X) = 0$인 것만 보이면 충분하다. (물론, $U$가 공집합이면 $1$)
  • $U$의 부분 multiset 중 중복 원소가 존재하는 것은 좌변에 기여하지 않는다. 그러니 좌변에 기여하는 것은 중복 원소가 없는 경우 뿐.
  • 그러니 단순히 $U$에서 중복 원소를 제거한 집합을 보고, 기존 결과를 그대로 적용하면 증명이 끝난다. 


이제 정수론에서의 mobius function을 이야기 할 수 있다. 

  • 자연수 $n$을 소인수분해하여, $p_1^{e_1} p_2^{e_2} \cdots p_k^{e_k}$라 쓰자. 이제 $n$을 multiset $\{p_1, \cdots , p_1, p_2, \cdots , p_2, \cdots p_k \}$로 대응시킨다. 단, $p_i$는 $e_i$번 등장.
  • 이제 "집합의 포함 관계"는 "자연수의 나누어떨어짐 관계"로 대응이 되며, mobius function 역시 그대로 대응된다. 
  • $\sum_{d|n} \mu(d) = [n=1]$ 역시 $\sum_{S \subset A} \mu(S) = [A = \emptyset] $과 그대로 대응된다.
  • 그러니 포함-배제와 mobius inversion, mobius function에 관한 문제풀이는 다를 것이 없다. 
  • 즉, 독자가 포함-배제에 관해서 가지고 있는 문제풀이 아이디어, 전략이나 팁들 대부분이 이번 7단원 내용에도 적용될 수 있다.


잠깐 정수론과 무관한 이야기.

  • $g(A)$의 값이 주어졌을 때 $f(A)$의 값을 전부 구하는 문제를 다시 생각해보자.
  • 순서대로 $f(A) = g(A) - \sum_{S \subsetneq A} f(S)$를 적용했다면, subset을 전부 iterate 해야하므로 $\mathcal{O}(3^n)$의 시간이 걸린다.
  • 하지만 $f(A) = \sum_{S \subseteq A} (-1)^{|A|-|S|} g(S)$임을 이용하면, $(-1)^{|A|} f(A) = \sum_{S \subseteq A} (-1)^{|S|} g(S)$임을 얻는다.
  • 이제 SOS DP를 이용하여 $f$ 전부를 $\mathcal{O}(2^n n)$ 시간에 계산할 수 있다. 시간 절약하기 좋은 테크닉.  


이제부터는 문제풀이에는 별로 중요하지 않은 이야기.

  • "나누어떨어짐 관계" $a|b$와 "포함 관계" $S \subset A$ 사이에는 무언가 비슷한 점이 있는 것 같다.
  • $a|a$이며, $a|b$, $b|a$면 $a=b$ (자연수 기준), $a|b$, $b|c$면 $a|c$가 성립한다.
  • $S \subset S$이며, $S \subset T$, $T \subset S$면 $S = T$, $S \subset T$, $T \subset U$면 $S \subset U$가 성립한다.
  • 이러한 조건을 만족시키는 relation을 partial order라고 한다. mobius function은 이러한 partial order를 갖는 poset에서 전부 정의할 수 있다.


관련 코드는 github에서 찾아볼 수 있다.


이 단원의 내용은 특정 값을 계산하는 알고리즘과 그 template code가 명확하게 있지 않다.

즉, 이 단원에 관한 문제 풀이는 template을 복사/붙여넣기하는 방식으로 풀기 어렵다. 

대신, 문제풀이를 위해서 알아야 하는 이론의 양은 상대적으로 적다. 적은 이론 + 많은 식 정리로 풀리는 문제가 대부분이다.


하지만 "이 단원의 내용이 포함-배제의 원리와 같다"는 점을 제대로 이해하면, 문제풀이가 더 쉬워질 수도 있다. 

그러나 이 점을 제대로 이해하지 않아도 문제는 풀 수 있으니, 이 부분은 부록으로 옮긴다. 그래도 부록을 한 번 읽는 것을 추천.


mobius function $\mu(n)$은 각 자연수에 대하여 다음과 같이 정의된다.

  • $n = 1$인 경우, $\mu(n) = 1$이다.
  • 만약 소수 $p$가 존재하여 $p^2|n$이라면, $\mu(n) = 0$.
  • $n = p_1p_2 \cdots p_k$라면 (단, $p_i$는 서로 다른 소수) $\mu(n) = (-1)^k$

이때, mobius function은 다음과 같은 성질을 갖는다.

  • $\mu$는 multiplicative. 즉, $m, n$이 서로소인 자연수라면 $\mu(mn) = \mu(m)\mu(n)$. 이는 정의에서 바로 증명할 수 있다.
  • 각 자연수 $n$에 대하여 $\sum_{d|n} \mu(d)$의 값은 $n=1$인 경우 $1$이고, $n \ge 2$인 경우에는 $0$이다. 증명은 뒤로 미룬다.
  • mobius inversion. $f(n) = \sum_{d|n} g(d)$라면, $g(n) = \sum_{d|n} \mu(n/d) f(d)$가 성립한다.
  • 예시. 오일러 정리 단원에서 $n = \sum_{d|n} \phi(d)$라고 했다. 그러니 $\phi(n) = \sum_{d|n}  d \mu(n/d)$다.
  • $\sum_{d|n} d \mu(n/d)$라는 식은, $n$의 소인수분해를 생각한 후 계산하면 $n(1-1/p_1)(1-1/p_2) \cdots (1-1/p_k)$와 같다.
  • $\mu(1), \mu(2), \cdots , \mu(n)$을 전부 계산하는 것은 에라토스테네스의 체로 할 수 있다. github 코드 참조.


사실 이게 CP/PS에서 쓰는 성질 전부다. 이걸 잘 조합해서 문제를 풀어내는 것이 어려운 것이다. 

  • mobius function이란 말을 들으면 mobius inversion을 생각하는 경우를 봤다. 하지만 실제로 mobius inversion 식 자체는 많이 쓰이지 않는다.

  • 사실 mobius function을 이용한 문제풀이의 핵심은 $\sum_{d|n} \mu(d)$가 $n=1$일 때만 $1$, 아니면 $0$이라는 점이다. 

  • 계속 강조하지만, mobius inversion은 포함-배제와 같다. 부록 참고.


이제부터 시범을 보인다. 식 조작의 편의를 위하여, Iverson bracket을 도입한다.

  • 명제 $P$에 대하여, $[P]$의 값은 $P$가 참이면 $1$, 아니면 $0$이다. 

  • 이제 $\sum_{d|n} \mu(d) = [n=1]$이라고 쓸 수 있다. 이 식이 문제풀이의 핵심이다.


BOJ 16409 : Coprime Integers


문제에서 요구하는 값은 $$\sum_{x=a}^b \sum_{y=c}^d [\text{gcd}(x, y)=1]$$이다. $\sum_{d|n} \mu(d) = [n=1]$을 사용하여 Iverson bracket을 없애면, $$ \sum_{x=a}^b \sum_{y=c}^d \sum_{g|x, g|y} \mu(g) $$를 얻는다. 이제 식을 바라보는 관점을 바꾸자. $x, y$에 대해서 iterate 한 다음 $g|x, g|y$를 만족하는 $g$를 보는 것이 아니라, $g$를 고정하고 $g$의 배수가 되는 $x, y$에 대해서 iterate 하자. 이렇게 식을 변형하는 방식을 보통 "double counting"이라고 부른다. 그러면 우리의 식이 $$ \sum_{g=1}^{\text{min}(b, d)} \mu(g) \cdot ([a, b] \text{ 에 존재하는 } g \text{의 배수 개수}) \cdot ([c, d] \text{ 에 존재하는 } g \text{의 배수 개수}) $$가 된다. 물론, $[a, b]$에 존재하는 $g$의 배수 개수는 $\lfloor b/g \rfloor - \lfloor (a-1)/g\rfloor$이고, $[c, d]$에 대해서도 마찬가지다. 즉, 답은 $$ \sum_{g=1}^{\min(b,d)} \mu(g) \cdot (\lfloor b/g \rfloor - \lfloor (a-1)/g \rfloor) \cdot (\lfloor d/g \rfloor - \lfloor (c-1)/g \rfloor)$$가 된다. 이제 $\mu$를 전처리하면 답을 쉽게 계산할 수 있다.


여기서도 볼 수 있지만, 약수에 대해 iterate 하는 것보다 배수에 대해 iterate 하는 것이 더 쉽다는 점을 많이 이용한다. (아래에 예시를 더 넣어두겠다)


시범을 한 번 더 보이겠다. 이번에는 $$\sum_{x=1}^n \sum_{y=1}^n \text{gcd}(x, y)$$를 구해보자. 이를 위해서 $\text{gcd}(x, y) = d$인 $(x, y)$의 개수를 각 $d$에 대해서 구해보자. 이 값을 구하려면, $$ \sum_{x=1}^n \sum_{y=1}^n [\text{gcd}(x, y)=d]$$를 구하면 된다. 그러니, 기존 식의 값은 $$ \sum_{d=1}^n d \cdot \sum_{x=1}^n \sum_{y=1}^n [\text{gcd}(x, y)=d]$$와 같다. $\text{gcd}(x, y)=d$려면, $x, y$ 역시 $d$의 배수여야 하니 $x=da$, $y=db$라 쓸 수 있고, 이때 $1 \le a, b \le \lfloor n/d \rfloor$와 $\text{gcd}(a, b)=1$이다. 즉, 식은 $$ \sum_{d=1}^n d \cdot \sum_{a=1}^{\lfloor n/d \rfloor} \sum_{b=1}^{\lfloor n/d \rfloor} [\text{gcd}(a, b)=1]$$과 같다. 이제 바로 앞의 문제의 결과를 적용하면 $$ \sum_{d=1}^n \sum_{g=1}^{\lfloor n/d \rfloor} d \mu(g) \cdot \lfloor \frac{n}{gd} \rfloor^2$$을 얻는다. 여기서 자연수 $n, a, b$에 대하여 $\lfloor n/(ab) \rfloor = \lfloor \lfloor n/a \rfloor / b \rfloor$임을 이용하였다.

  • Finish 1. $\mu$를 전처리하자. 위 합을 for 문으로 계산하면, 시간복잡도는 $\sum_{d=1}^n \lfloor n/d \rfloor = \mathcal{O}(n \log n)$ (by Harmonic Sequence)

  • Finish 2. $gd$라는 값이 심상치 않다. $gd = T$라고 치환을 한 다음, $T$를 기준으로 계산을 해보자. 그러면 $$ \sum_{T=1}^n \lfloor n/T \rfloor^2 \cdot \sum_{d|T} (d \cdot \mu(T/d))$$라는 식을 얻는다. 그런데 $\phi(T) = \sum_{d|T} (d \cdot \mu(T/d))$임을 앞에서 언급했다. 그러니 식은 $ \sum_{T=1}^n \lfloor n/T \rfloor^2 \phi(T)$와 같다. 

위 문제의 mobius function을 사용하지 않는 풀이를 소개한다. $cnt[d]$를 $\text{gcd}(x, y) = d$인 $x, y$의 개수라고 하자. 사실 우리는 $\text{gcd}(x, y) = d$인 $x, y$의 개수는 세기 어려워해도, $\text{gcd}(x, y)$가 $d$의 배수인 것은 쉽게 셀 수 있다. $x, y$가 각각 $d$의 배수이면 충분하므로, 정확히 $\lfloor n/d \rfloor^2$개 존재할 것이다. 이는 사실 $$ cnt[d] + cnt[2d] + \cdots + cnt[d \cdot \lfloor n/d \rfloor]  = \lfloor n/d \rfloor^2$$이 성립한다는 것이다. 이를 통해서, $cnt$ 배열을 모두 구할 수 있다. 다음과 같은 코드를 생각하면 이해가 될 것이다.


1
2
3
for(i=1 ; i<=n ; i++) cnt[i] = (n/i) * (n/i);
for(i=n ; i>=1 ; i--)
    for(j=2*i ; j<=n ; j+=i) cnt[i]-=cnt[j];
cs


더욱 놀라운 점은, Harmonic Sequence의 성질 때문에 이 코드 역시 시간복잡도가 $\mathcal{O}(n \log n)$이란 것이다.

  • 이 풀이를 보면서, "배수에 대해 iterate 하는 것이 효율적이다"라는 말의 의미를 느낄 수 있다.

  • 또한, 이 풀이가 굉장히 포함-배제의 원리 느낌이 난다는 것을 느낄 수 있을 것이다. 

물론, 실제로 문제를 해결하려면 $\sum d \cdot cnt[d]$를 계산하면 끝이다. 


이제부터 "배수에 대해 iterate 하는 것"과 double counting을 사용한 간단하지만 중요한 예제를 몇 개 제시한다.

  • $\tau(n)$을 $n$의 약수의 개수라고 정의한다. $\sum_{i=1}^n \tau(i)$를 어떻게 계산할까?  

  • $\sigma(n)$을 $n$의 약수의 합이라고 정의한다. $\sum_{i=1}^n \sigma(i)$를 어떻게 계산할까?

$\sum_{i=1}^n \tau(i)$는 결국 $g|i$, $1 \le i \le n$을 만족하는 순서쌍 $(g, i)$의 개수이다. 순서쌍의 개수를 $i$ 기준으로 세면, 각 $i$에 대해 조건을 만족하는 $g$가 $\tau(i)$개 있기 때문이다. 순서쌍의 개수를 $g$ 기준으로 세면, 조건을 만족하는 $i$가 $\lfloor n/g \rfloor$개 있음을 알 수 있다. 그러니 순서쌍의 개수는 $\sum_{g=1}^n \lfloor n/g \rfloor$이다. 그러므로, $$ \sum_{i=1}^n \tau(i) = \sum_{i=1}^n \lfloor n/i \rfloor $$이다. 이는 결국 "각 정수에 대해 약수의 개수를 센다"에서 "각 정수에 대해 배수의 개수를 센다"로 문제를 바꿔 생각해서 얻은 결과이다. 위 식의 우변은 자명하게 $\mathcal{O}(n)$에 계산할 수 있다. 또한, $\lfloor n/i \rfloor$로 가능한 값의 개수가 $\mathcal{O}(\sqrt{n})$개이며, 그 값을 갖게 하는 $i$의 범위 역시 계산할 수 있음은 이미 1단원에서 언급하였다. 이를 사용하면, 우변의 값을 $\mathcal{O}(\sqrt{n})$에 계산할 수 있다. 이러한 방식은 후에 multiplicative function을 다룰 때 더 나온다.


$\sum_{i=1}^n \sigma(i)$ 역시, 같은 방법으로 생각하면 $\sum_{g=1}^n g \lfloor n/g \rfloor$임을 얻는다. 이는 $\mathcal{O}(n)$에 자명하게 계산할 수 있다. 특정 구간에 속한 정수의 합, 즉 $\sum_{i=a}^b i$를 쉽게 계산할 수 있으므로, $\lfloor n/g \rfloor$로 가능한 값이 $\mathcal{O}(\sqrt{n})$개임을 이용하여 우변의 값을 $\mathcal{O}(\sqrt{n})$에 계산할 수 있다. github 코드 참고.  

관련 코드는 github에서 찾아볼 수 있다.


이 글에서 다룰 내용의 많은 부분에서 완벽한 수학적 논의는 생략할 것이다.


Miller-Rabin 소수 판별법


이 부분에 대해서는 이미 좋은 자료가 있습니다. https://casterian.net/archives/396


우리는 주어진 $n$이 소수임을 확인하고자 한다. 먼저 $n-1 = d \cdot 2^s$인 정수 $s, d$를 찾는다.

만약 $n$이 소수라면, 각 자연수 $1 \le a < n$에 대하여 다음이 성립한다.

  • [P] $a^d \equiv 1 \pmod{n}$이거나, $r=0, 1, \cdots , s-1$ 중 적어도 하나에 대해 $a^{2^rd} \equiv -1 \pmod{n}$

증명을 해보자. $a^{n-1}-1 = a^{d \cdot 2^s}-1 = (a^d-1)(a^d+1)(a^{2d}+1)(a^{2^2d}+1) \cdots (a^{2^{s-1}d}+1)$이다.

그런데 페르마 소정리에 의하여 좌변은 $0 \pmod{n}$이므로, 우변 역시 $n$의 배수임을 알 수 있다.

$n$이 소수이므로, 여러 자연수의 곱이 $n$의 배수라는 것은 그 중 하나 이상이 $n$의 배수임을 의미한다. 

그러므로, $a^d \equiv 1 \pmod{n}$이거나 적당한 $0 \le r <s$가 있어 $a^{2^rd} + 1 \equiv 0 \pmod{n}$이 성립한다.


$n$이 주어졌을 때, $s, d$의 값을 계산하는 것은 쉽게 할 수 있다. 

또한, $1 \le a < n$을 만족하는 $a$가 주어졌을 때, [P]가 성립하는지 여부도 역시 빠르게 확인할 수 있다. 


이제 다음과 같은 알고리즘을 하나 고안할 수 있다.

  • 테스트 횟수 $k$를 하나 정한다. $k$를 어떻게 잡는가에 대한 이야기는 나중에 한다.
  • $1 \le a < n$을 만족하는 $a$를 랜덤하게 하나 잡자.
  • [P]가 성립하는지 확인하자. 만약 성립하지 않는다면, $n$이 소수가 아니라는 뜻이고 알고리즘 종료.
  • [P]가 성립한다면, 다시 처음으로 돌아가서 알고리즘을 반복한다. $k$번에 걸쳐 항상 [P]가 성립한다면 소수라고 판단하고 종료.

만약 $n$이 소수라면, $a$를 어떻게 잡더라도 [P]가 성립할 것이니, 무조건 소수라고 판단할 것이다.

$n$이 합성수라면 어떨까? $n$이 소수라면 [P]가 성립한다고 했지, $n$이 합성수면 무조건 [P]가 거짓이라고 하지는 않았다.


위 알고리즘이 정당한 소수 판별 알고리즘이 되려면, $n$이 합성수일 경우 높을 확률로 $n$이 소수가 아니라는 결론을 내려야 한다.

즉, $k$번의 테스트를 거쳤을 때 [P]가 거짓이 되는 $a$의 값이 매우 높은 확률로 한 번 이상 뽑혀야 한다. 이는 결국

  • $n$이 합성수인 경우, [P]가 거짓이 되는 $a$가 얼마나 많은가? 

라는 문제와 동치이다. 이제 중요한 결론을 증명하지 않고 제시한다. 증명은 구글 ㄱㄱ

  • $n$이 합성수인 경우, [P]가 거짓이 되는 $a$는 적어도 $3n/4$개 존재한다.
  • 즉, $n$이 합성수인 경우, 한 번의 테스트를 통과할 확률은 $1/4$ 이하이다. 

이는 $k$번의 독립적인 테스트를 거쳤을 때, 합성수가 모든 테스트를 통과할 확률이 $4^{-k}$ 이하라는 뜻이다. 

그러니 $k=20$ 정도로 잡아주면 매우 높은 확률로 정답을 반환하는 좋은 소수 판별 알고리즘이 될 수 있다.


하지만 랜덤이 싫을 수 있다. 랜덤을 제거한 Miller-Rabin 알고리즘의 여러 variant가 존재한다.

  • Generalized Riemann Hypothesis를 가정하면, $2 \le a < 2(\ln n)^2$를 전부 시도하면 정확하게 소수 판별을 할 수 있다.
  • $n < 2^{32}$라면, $a = 2, 7, 61$인 경우만 따져보면 된다. 이때, $n = 2, 7, 61$인 경우를 따로 처리해야 함에 주의하라.
  • $n < 2^{64}$라면, $a$로 $2$ 이상 $40$ 이하의 소수만 따져보면 된다. 이때, $n$ 역시 $2$ 이상 $40$ 이하의 소수인 경우에 주의하라.    

랜덤을 제거한 variant가 아마 실제 성능도 좋을 것이다. $n < 2^{64}$의 variant에서도 테스트를 12번만 하기 때문이다.


구현 노트

  • 구현이 까다로운 알고리즘은 아니지만, C++의 int 범위를 넘어가는 modulus에서 여러 연산을 해야한다. 
  • 특히 까다로운 것은 곱셈이다. $a, b, n$이 모두 64bit 정수일 때, $ab \pmod{n}$을 계산해야 한다.
  • 이를 위해서 쓸 수 있는 첫 번째 방법은 이집트 곱셈, 즉 이진법을 사용한 곱셈이다. 이를 위해서 로그 시간이 필요하다.
  • 두 번째 방법은 __int128을 동원하는 것이다. $ab$의 값을 __int128에서 계산한 뒤, $n$에 대한 나머지를 구한 후 64bit 정수로 반환한다. 
  • $n$이 작으면, $\mathcal{O}(\sqrt{n})$ 알고리즘이 Miller-Rabin 알고리즘보다 빠를 수 있다. 최적화가 필요한 경우 참고하자.
  • 더 적은 테스트 횟수로 $2^{64}$ 미만 자연수의 소수 여부를 판별할 수 있다. 이때, $a$ 값이 소수가 아닌 경우 "$a$의 약수인 소수"들을 특별히 고려해야 한다. 
  • 위 두 가지 포인트를 추가로 고려하여 작성한 Miller-Rabin 구현체를 원한다면, 맨 위에 소개한 casterian님의 블로그를 참고하라.


Pollard-Rho 소인수분해


이 부분 역시 좋은 자료가 있습니다. https://aruz.tistory.com/140


핵심 아이디어 1. Floyd's Cycle Detection Algorithm


잠깐 전형적인 CP/PS 토픽인 그래프로 넘어오자. functional graph는 각 정점이 outdegree 1을 갖는 directed graph다.

정점 $v$가 있으면 그 정점에서 나가는 간선이 정확히 하나이니, 그 간선을 따라 도달하는 정점을 $f(v)$라 정의할 수 있다. 

즉, $f(v)$는 $v$의 "다음 정점"이다. 이제 $v \rightarrow f(v) \rightarrow f(f(v)) \rightarrow \cdots$라는 path를 생각하자.

결국 이 경로는 하나의 일직선을 그리다가, 이미 방문한 정점을 다시 방문하면서 cycle에 들어가게 될 것이다. 이 cycle을 어떻게 찾을까?


물론 여러 방법이 있겠으나, Floyd의 "토끼와 거북이" 알고리즘이 매우 효율적이다. 

  • 잠시 $f^n(v) = f(f^{n-1}(v))$로 정의하자. 즉, $f^n(v)$는 $v$의 "$n$번째 뒤"의 정점이다.
  • 토끼 한 마리와 거북이 한 마리를 $v$에 두자. 각 동물들의 위치를 변수 $r, t$로 명명하고 관리하겠다.
  • 각 단계마다, 토끼는 두 칸, 거북이는 한 칸 이동시킨다. 즉, $r \leftarrow f(f(r))$, $t \leftarrow f(t)$를 반복한다.
  • 만약 $r = t$인 시점이 오면, 사이클을 찾은 것이다. 물론, 맨 처음에 $r=t=v$인 경우는 제외하고 생각한다.
  • 왜 사이클을 찾은 것인가? $k$번의 단계를 거친 후에 $r = t$임을 확인했다고 가정하자.
  • 이는 $f^k(v) = f^{2k}(v)$임을 의미한다. 즉, $k$번째 뒤 정점과 $2k$번째 뒤 정점이 같다는 것을 의미한다.
  • 이제 $k$번째 정점인 $f^k(v)$부터 시작하여, 다음 정점으로 이동하는 것을 반복한다. 
  • 다시 시작한 $f^k(v)$에 도착할 때까지 이를 반복하면, 지금까지의 경로가 완벽하게 찾고자 하는 사이클이 된다.
  • 왜 사이클이 찾아지는가? $i < j$이고 $f^i(v) = f^j(v)$인 $i, j$가 존재할 때, (즉 사이클이 있을 때) 자연수 $k$가 있어 $f^k(v) = f^{2k}(v)$임을 보이면 된다.
  • "한 칸씩 격차가 벌어지는 걸 결국 사이클 길이만큼 격차가 벌어져서 다시 만났다"고 생각하면 직관적이며, 증명도 똑같다.
  • 속도 차이가 나는 두 사람이 400m 트랙에서 달렸을 때, 실제로 달린 거리가 400m 차이가 나서 만날 수 있는 것과 같은 원리.
  • 자세한 증명을 원한다면, $t(j-i) \ge i$인 자연수 $t$를 하나 잡고, $k=t(j-i)$라고 정의해보면 증명됨을 확인할 수 있다.
  • 왜 효율적인가? 일단 시간복잡도가 (사이클 도달에 필요한 거리) + (사이클의 길이) 정도로 매우 효율적이다.
  • 공간복잡도도 효율적이다. 각 정점 $v$에 대해서 $f(v)$를 계산하는데 드는 메모리를 제외하면, $\mathcal{O}(1)$만큼의 메모리가 추가적으로 필요하다.
  • 특히, 여러 경우에서 $f(v)$는 진짜 간선으로 표현되는 "다음 정점"이 아니라, 말 그대로 $v$에 대한 함수가 된다. 
  • 이는 아래의 경우도 마찬가지. 애초에 이름이 functional graph인 이유가 있다. 이 경우, $f(v)$를 계산하는데 드는 메모리도 $\mathcal{O}(1)$.


핵심 아이디어 2. Pseudorandom Sequence. Birthday Paradox


Birthday Paradox는 $\{1, 2, \cdots , n\}$에서 $\sqrt{n}$급 개수의 원소를 독립적으로 뽑으면, 같은 원소를 두 번 이상 뽑을 확률이 상당히 높다는 결과이다.

이는 $n=365$일 때, 대충 20~30명의 학급에서 생일이 겹치는 경우가 많다는 내용으로 알려진 결과일 것이다.


이제 우리의 목표는 주어진 합성수 $n$을 소인수분해 하는 것이다. $p$를 $n$의 가장 작은 소인수라고 하자. 우리는 $n$은 알고 있으나 $p$는 모른다.


이제 다음과 같은 방식으로 랜덤한 수열을 잡는다. $x, c$ 각각을 $1$ 이상 $n$ 미만 랜덤한 자연수라고 두자.

$x_0 = x$, $x_i = (x_{i-1}^2 + c) \pmod{n}$으로 수열을 잡자. 여기서 짚고 넘어가야 할 부분은

  • 수열 $\langle x_k \rangle$은 pseudorandom. 즉, 완전 랜덤한 것은 아니지만 랜덤에 가깝다.
  • $x_i$의 값은 $x_{i-1}$의 값에 의하여 완전히 결정된다. 그러니, $x_{i-1} \rightarrow x_i$를 간선으로 생각하는 functional graph를 생각할 수 있다.
  • 즉, 이 수열의 값을 정점으로 갖는 functional graph를 생각할 수 있고, 마찬가지로 cycle 개념도 생각할 수 있다.
  • $x_i \equiv (x_{i-1}^2 + c) \pmod{p}$ 역시 성립한다. 그러니, $y_i = x_i \pmod{p}$로 정의하면, $y_i \equiv (y_{i-1}^2 + c) \pmod{p}$이다.
  • 현재 계산할 수 있는 것은 $x_i$이며, $y_i$는 실제로는 $p$를 모르기 때문에 계산할 수 없는 값이다.

이제 본격적으로 알고리즘을 설계해보자.

  • $\text{gcd}(x_{2k} - x_k, n)$이라는 값을 생각해보자. 이 값을 계산하려면 $x_{2k}, x_k$의 값이 필요하며, 이는 "토끼와 거북이"의 방식으로 관리할 수 있다.
  • $y_k = y_{2k}$라면, (즉 수열 $y$에서 cycle이 발견되었다면) $\text{gcd}(x_{2k} - x_k, n)$은 $p$의 배수가 된다. 
  • Birthday Paradox에 의해, 이렇게 수열 $y$에서 cycle이 발견되려면 평균적으로 $\mathcal{O}(\sqrt{p})$ 정도 index가 되어야 한다. 
  • 비슷하게, 수열 $x$에 대해서 cycle이 발견되려면 평균적으로 $\mathcal{O}(\sqrt{n})$ 정도 index가 되어야 한다.
  • 즉, 평균적으로 $\mathcal{O}(\sqrt{p})$ 정도의 계산량 이후로 $\text{gcd}(x_{2k}-x_k, n)$이 $1$이 아닌 시점이 오며, 이때 $\text{gcd}(x_{2k}-x_k , n)$은 높은 확률로 $n$이 아니다.
  • 이는 $\text{gcd}(x_{2k}-x_k, n)$이 $n$의 "nontrivial divisor"가 된다는 뜻이다. 이를 반복하면 소인수분해를 얻는다.

$n$의 non-trivial divisor를 얻기까지 평균적으로 $\mathcal{O}(\sqrt{p}) = \mathcal{O}(n^{1/4})$번의 GCD 연산이 필요하다.


구현 노트.

  • 이 알고리즘에서는 random을 사용한다. C++의 rand()를 사용하기 보다는, mt19937 등을 사용하는 것을 추천한다.
  • 앞에서도 언급했지만 $n$이 합성수라는 조건이 붙는다. 미리 Miller-Rabin 소수 판별을 통해 $n$이 소수인지 판별을 해야한다.
  • $\text{gcd}(x_{2k}-x_k, n)=n$일 확률이 낮다고 했지, 불가능하다고 하지는 않았다. 이 경우, 시작하는 값 $x, c$를 다시 잡고 알고리즘을 다시 돌리자.
  • $\text{gcd}(x_{2k}-x_k, n)=d$가 $1, n$이 아니라면, 문제를 $d$와 $n/d$를 각각 소인수분해 하는 것으로 바꿀 수 있다. 이제 재귀적으로 문제를 풀자.
  • 앞에서 언급한 64bit 정수의 곱셈에 대한 내용이 여기서도 적용된다. 이집트 곱셈을 쓰거나, __int128을 사용하자.


2020년은 코로나 바이러스 때문에 모두에게 힘든 한 해였을 것으로 생각합니다. 저도 처음에는 괜찮았는데 만나고 싶은 사람도 계속 못 만나고 먹고 싶은 음식도 제대로 못 먹으니까 정신적으로 지치네요. 하지만 주변에 좋은 사람을 많이 둔 덕분에, 2020년은 놀랍게도 제가 살면서 가장 많이 성장한 1년이 될 수 있었습니다. 예전에는 (2016~2019) 학점에 과하다 싶을 정도로 투자해서 상대적으로 덜 중요한 교양 과목에도 많은 시간을 넣었고, 그 때문에 학점을 챙기는 것과 PS를 제외하면 특별히 하고 있는 일이 없었습니다. 흥미로운 전공도 극히 일부를 제외하면 듣지 못했습니다. 그에 비해, 올해는 수업도 유용하고 재밌는 것을 많이 들었고, 학점도 적당히 챙겨주면서 학교 밖에서도 많은 것을 배우고 이룰 수 있었습니다. 이렇게 매년 성장할 수 있으면 정말 좋을 것 같네요.


이제부터 글은 토픽으로 나누어서 진행됩니다. 사진 대부분은 링크로 대체됩니다. 


PS/CP 이야기


대학 이야기


CTF/암호학 이야기


일상 이야기


미래 이야기


'미래 계획 및 근황' 카테고리의 다른 글

2월 정리  (0) 2021.02.28
1월 정산 + 2월 계획 + 1학기 수강신청  (1) 2021.02.01
1학기 후기 || 최근에 한/할 것들 || 2학기 계획  (0) 2020.06.23
여름방학 목표  (0) 2020.06.19
겨울방학 목표 정리  (4) 2019.11.29

관련 코드는 github에서 찾아볼 수 있다.


이 글에서는 팩토리얼과 이항계수를 $n$으로 나눈 나머지를 계산하는 방법에 대하여 다룬다.

  • 중국인의 나머지 정리를 기반으로 한 아이디어는 여기서도 등장한다.
  • $n = p_1^{e_1} p_2^{e_2} \cdots p_k^{e_k}$에 대한 나머지를 구하려면, 각 $p_i^{e_i}$에 대한 나머지를 구하면 된다.
  • 그 후, 단순히 중국인의 나머지 정리를 사용하여 합쳐주면 된다. 그러니 문제를 prime power에 대해서만 해결하면 된다.

이 글에서 다룰 내용은 상당한 난이도를 갖는 강화판들이 있다. 이에 대한 것은 결과만 언급하고 링크를 걸겠다.


먼저 팩토리얼, 이항계수를 $p$로 나눈 나머지를 계산하는 방법을 알아보자. $p$는 소수.


팩토리얼


주제 1. 주어진 $n, p$에 대하여 $n! \pmod{p}$ 계산하기. 

단순히 $n! \pmod{p}$를 계산한다고 가정하자. $n \ge p$인 경우, 값은 $0$이다. 

그러니, $n < p$인 경우만 고려해도 충분하다. $n$이 작으면, $\mathcal{O}(n)$ 알고리즘이 자명하다.

$n, p$가 모두 $10^{9}$-scale로 큰 경우는, $\mathcal{O}(\sqrt{p}\log p)$ 시간에 해결할 수 있으나 상당히 어려운 문제다. 문제와 이 글 참고. 


주제 2. $1!, 2!, \cdots , n!$ 각각을 $p$로 나눈 나머지와 각각의 $\pmod{p}$에 대한 modular inverse 계산하기.

modular inverse를 계산하는 이유는 역시 이항계수를 계산하기 위해서이다. 이 주제는 수많은 카운팅 문제들에서 사용된다.

$1!, 2!, \cdots n!$ 각각을 $p$로 나눈 나머지는 점화식 $k! \equiv (k \cdot (k-1)!) \pmod{p}$를 이용하여 $\mathcal{O}(n)$ 시간에 계산할 수 있다.


팩토리얼들의 modular inverse를 구하는 것은 정말 여러 방법이 있다.

  • 방법 1. 우선 각각의 modular inverse를 그냥 구하는 방법이 있다. 확장 유클리드 알고리즘을 쓰면 된다. 이 경우 $\mathcal{O}(n \log p)$의 시간 소요.
  • 방법 2. $1, 2, \cdots, n$ 각각의 modular inverse를 $\mathcal{O}(n)$에 구하는 방법이 있다. 이를 사용하면 팩토리얼의 modular inverse도 똑같이 $\mathcal{O}(n)$에 ok.
  • 알고리즘의 설명을 위해 프로그래밍 같은 notation을 잠시 도입한다. $i$의 modular inverse를 $inv[i]$라 하고, $a$를 $b$로 나눈 나머지를 $a \% b$라 하자.
  • Fact. $2 \le i<p$이면, $inv[i] \equiv inv[p \% i] \cdot (p - \lfloor p/i \rfloor) \pmod{p}$가 성립한다. ($p$가 소수임이 필요하다)
  • 증명. $p = qi + r$이라 하자. $p \% i = r$이며, $\lfloor p/i \rfloor = q$이다.
  • 이제 $i \cdot inv[p \% i] \cdot (p - \lfloor p/i \rfloor) \equiv i \cdot inv[r] \cdot (p-q) \equiv inv[r] \cdot (-qi) \equiv inv[r] \cdot r \equiv 1 \pmod{p}$이므로 증명 끝.
  • 여기서 $p \% i < i$이며, $inv[1]=1$임을 생각하자. 위 Fact는 결국 $inv$ 배열을 채우기 위한 DP 식이 된다. 
  • $k!$의 modular inverse는 $(k-1)!$의 modular inverse와 $k$의 modular inverse를 곱해서 얻을 수 있다.
  • 방법 3. 매우 쉽고 똑똑한 아이디어다. $n!$의 modular inverse를 먼저 확장 유클리드 알고리즘으로 구하자.
  • 그 후, $k!$의 modular inverse가 $(k+1)!$의 modular inverse와 $(k+1)$의 곱임을 이용한다. 역순으로 계산하는 것이다.
  • 이 방법으로 $1, 2, \cdots n$ 각각의 modular inverse도 $\mathcal{O}(n)$ 시간에 구할 수 있다. 
  • $k$의 modular inverse를 구하려면, $k!$의 modular inverse와 $(k-1)!$을 곱하면 된다.

방법 2, 3은 둘 다 $\mathcal{O}(n)$ 알고리즘이다. 하나만 알아도 충분하지만, 극한의 시간 최적화가 필요한 경우에는 둘 다 알아야 한다.

  • 만약 $1, 2, \cdots , n$ 각각의 modular inverse를 얻는 것이 주요 목적이라면, 방법 2를 사용하는 것이 맞다.
  • 만약 $1!, 2!, \cdots, n!$ 각각의 modular inverse를 얻는 것이 주요 목적이라면, 방법 3을 사용하는 것이 맞다.
  • 위에서 언급한 modular inverse들을 전부 구하는 것이 목적이라면, 어느 것을 사용해도 문제가 없다.
  • 이유를 파악하기 위해서는 각 방법들이 언급한 값들을 계산하기 위해서 몇 번의 나머지 연산을 하는지 생각해보자.

개인적인 추천은 그냥 방법 2, 3을 다 알고 있는 것이다. 방법 2의 원리까지는 이해할 필요가 없으니, 코드만 알아두면 된다.


주제 3. $n! = p^e m$라고 쓸 경우 (단, $m$은 $p$의 배수가 아니다) $e$와 $m \pmod{p}$를 계산하기.

이는 일종의 "$n!$의 0이 아닌 마지막 자리 구하기"와 같은 부류의 문제라고 볼 수 있다. 이 문제를 $p$진법에서 푸는 것이기 때문.

이 문제가 필요한 이유는 이항계수의 계산이다. $n!$이 $p$의 배수라고 해서 $\binom{n}{k} =n!/(k!(n-k)!)$이 $p$의 배수일 이유는 없다.


$n < p$인 경우는 $e=0$이며, $m \pmod{p}$는 단순히 $n! \pmod{p}$이다. 그러니 앞선 논의를 그대로 사용할 수 있다.

사실 이 말은 주제 3의 문제가 주제 1의 문제를 포함하고, 더 어려운 문제라는 이야기다. 그러니 $p$가 크면 이 문제도 풀기 힘들 것이라고 예측할 수 있다.

여기서 다룰 경우는 $p$가 $\mathcal{O}(p)$ 알고리즘이 돌아갈 정도로 작고, $n$은 매우 클 수 있는 경우이다.


여기서도 문제를 축소하는 방법을 사용할 것이다. $n=pk+r$, $0 \le r < p$라고 하자. 아래 그림과 함께 보자.

  • $1$부터 $n$까지의 수 중, $p$의 배수가 아닌 것에 먼저 집중하자.
  • $[1, p-1], [p+1, 2p-1], \cdots , [p(k-1)+1, pk-1], [pk+1, pk+r]$ 형태의 구간으로 나누어 볼 수 있을 것이다.
  • 첫 $k$개의 구간을 보자. 각 구간에 있는 정수들을 곱한 값은 $(p-1)! \pmod{p}$이다. (사실 이는 $-1 \pmod{p}$다. 윌슨의 정리)
  • $[pk+1, pk+r]$의 구간에 있는 정수들을 곱한 값은 $r! \pmod{p}$이다. 
  • 이제 $p$의 배수들인 것에 집중하자. $p, 2p, \cdots , kp$가 있다. 여기서 $p$를 $k$개 뽑아낼 수 있다.
  • 그러면 남는 것은 $1, 2, \cdots , k$고, 이는 정확히 $k!$과 같다. 여기서 우리는 문제를 $k!$에 대한 문제로 축소시켰다.
  • $k!$을 $p^{e'} m'$으로 쓸 때, $e'$와 $m' \pmod{p}$의 값을 알아냈다고 가정하자. 이제 $n! = p^e m$이라 쓸 때 $e$와 $m \pmod{p}$를 구할 수 있을까?
  • 가능하다. 위 논의를 정리하면, $e = e' + k$, $m \equiv m' \cdot (p-1)!^k \cdot r! \pmod{p}$를 반환하면 된다.
  • 앞에서도 언급했지만 $(p-1)! \equiv -1 \pmod{p}$이므로, $m$을 계산할 때는 $k$의 홀짝성만 봐도 충분하다.
  • 당연히 구현은 재귀적으로 하는 것이 편하다. 또한, 알고리즘을 돌리기 전 $1!, 2!, \cdots , (p-1)! \pmod{p}$를 전처리하는 것이 좋다. 
  • 문제가 얼마나 축소되는가? : $n$에서 $k$로 축소시켰고, $k \le n/p \le n/2$이다. 그러니 반 이상 감소하고, 로그 시간에 ok.
  • 여러분도 눈치챘겠지만, 여기서 $e = \lfloor n/p \rfloor + \lfloor n/p^2 \rfloor + \cdots $임을 얻는다.



이항계수


이제 이항계수 $\binom{n}{k}$를 $\pmod{p}$에 대하여 계산하는 방법을 보자. $\binom{n}{k} = \binom{n}{n-k}$이니, $k \le n/2$만 보자. 여러 경우가 있다.


Case 1. $k<p$이며 $k$가 작은 경우. 

이 경우, $\displaystyle \binom{n}{k} = \frac{n \cdot (n-1) \cdots (n-k+1)}{1 \cdot 2 \cdots k}$임을 이용하여 직접 계산을 할 수 있다.

분자, 분모를 모두 $\pmod{p}$로 구한 후, 분모의 modular inverse를 구해 나눗셈을 곱셈으로 대체한다. 

물론, $k<p$이므로 $\text{gcd}(k!,p)=1$이며, 그러니 확장 유클리드 알고리즘으로 modular inverse를 구할 수 있다.

사실, $\displaystyle \binom{n}{k} = \frac{n-k+1}{k} \binom{n}{k-1}$이 성립함을 이용하여 더 많은 것을 할 수 있다.

위 식을 차례대로 적용하여, 단순히 $\binom{n}{k} \pmod{p}$를 구하는 것이 아니라 $\binom{n}{0}, \binom{n}{1}, \cdots , \binom{n}{k}$ 각각을 $p$로 나눈 나머지를 구할 수 있다.


Case 2. 작은 $n, k$에 대하여 전부 전처리를 하고 싶은 경우.

이 경우, 파스칼의 삼각형을 이용하여 DP와 같은 방식으로 계산을 할 수 있다. 정수론이 필요없는 영역이다.


Case 3. $p$가 작은 경우, 즉 $\mathcal{O}(p)$ 알고리즘이 동작할 수 있는 경우.

먼저 $1!, 2!, \cdots , (p-1)!$ 각각을 $p$로 나눈 나머지와, 그들의 modular inverse를 전처리하자.


방법 1. Lucas 정리 사용

많이 사용되는 방법 중 하나다. 다음 결과를 증명하지 않고 제시한다. 증명은 위 링크에 걸려있다. 

정리. $m = m_km_{k-1} \cdots m_0$, $n=n_k n_{k-1} \cdots n_0$가 $m, n$의 $p$진법 전개라면, $\displaystyle \binom{m}{n} \equiv \binom{m_k}{n_k} \cdot \binom{m_{k-1}}{n_{k-1}} \cdots \binom{m_0}{n_0} \pmod{p}$다.

증명의 아이디어가 (Frobenius Endomorphism) 문제 풀이에 사용된 걸 본 적이 없는 것은 아니나, 모두 매우 어려운 문제들이다.

이 결과는 문제를 단순히 $m, n$이 $p$ 미만일 경우에 $\binom{m}{n} \pmod{p}$를 계산하는 문제로 만들어준다.

$m<n$인 경우 결과는 $0$이며, 그렇지 않은 경우 $m!/(n!(m-n)!)$을 계산하여 답을 얻을 수 있다.

이제는 어느 정도 당연한 이야기지만, 나눗셈은 modular inverse를 곱하는 것으로 계산한다.

전처리를 했기 때문에, 각 $\binom{m_i}{n_i} \pmod{p}$를 구하는 것에는 $\mathcal{O}(1)$의 시간이 필요하다.

그러니 시간복잡도는 $\mathcal{O}(k) = \mathcal{O}(\log_p m)$이 되겠다. 매우 빠르고 간단한 방법이다. 


방법 2. 앞서 다룬 주제 3 사용

주제 3을 이용하여도 문제를 해결할 수 있다. $\displaystyle \binom{n}{k} = \frac{n!}{k!(n-k)!}$이라 하자.

이제 $n! = p^{e_1} m_1$, $k! = p^{e_2} m_2$, $(n-k)!=p^{e_3} m_3$라고 쓰자. $m_1, m_2, m_3$는 $p$의 배수가 아니다.

이때, 주제 3에서 다룬 알고리즘으로 $e_1, e_2, e_3$와 $m_1, m_2, m_3 \pmod{p}$를 로그 시간에 계산할 수 있다.

그러면 이제 $\displaystyle \binom{n}{k} =  p^e m$이라 쓸 때, $e$와 $m \pmod{p}$를 구할 수 있을까?

있다. $e = e_1 - e_2 - e_3$, $m \equiv m_1 (m_2m_3)^{-1} \pmod{p}$가 성립하기 때문이다. $a^{-1}$은 $a$의 잉여역수.

만약 $e>0$이라면 주어진 이항계수는 $p$의 배수가 된다. 그렇지 않다면, 이항계수를 $p$로 나눈 나머지는 $m \pmod{p}$와 같다. 끝.


이제 $\pmod{p}$의 논의를 끝냈다. 이제 본격적으로 prime power에 대하여 논의해보자. 


팩토리얼


주제 1. 주어진 $n, p$에 대하여 $n! \pmod{p^e}$ 계산하기.

$n \ge pe$이면 $n!$에서 $p$의 배수가 $e$개 이상 들어가므로, $n! \pmod{p^e}$는 $0$이 된다. 

$n$이 작으면 단순한 $\mathcal{O}(n)$ 알고리즘이 작동한다는 점은 기존 주제 1의 논의에서와 동일하다. 


주제 2. $1!, 2!, \cdots , n!$ 각각을 $p^e$로 나눈 나머지와 각각의 $\pmod{p^e}$에 대한 modular inverse 계산하기.

문제가 잘 정의되려면 역시 $n < p$를 가정해야 한다. 이 경우, 방법 1과 방법 3이 잘 적용된다. 


주제 3. $n! = p^t m$라고 쓸 경우 (단, $m$은 $p$의 배수가 아니다) $t$와 $m \pmod{p^e}$를 계산하기.

$t$를 계산하는 것은 기존 주제 3과 애초에 목표가 동일하다. $m \pmod{p^e}$를 계산하는 것이 목표. 이 역시 기존 주제 3과 비슷한 방법으로 할 수 있다.

  • 쉽게 가자. $n=pk+r$, $0 \le r < p$라고 쓰자. 
  • $1$부터 $n$까지의 수 중, 이번에는 $p$의 배수인 것에 먼저 집중해보자.
  • 여기서는 $p$가 $k$개 추가되고, $k!$이 튀어나온다. 이는 문제의 "축소"된 형태이다. 여기까진 기존 주제 3과 동일.
  • 이제 $p$의 배수가 아닌 것을 봐야 한다. 잠시 $n=p^ek' + r'$이라고 하자. $0 \le r' < p^e$. 
  • $\mathcal{O}(p^e)$의 시간을 투자하여, DP를 이용해 다음 배열 $val[i]$를 ($0 \le i < p^e$) 채워놓자. 미리 전처리하는 것이다.
  • $val[i]$ : $1$ 이상 $i$ 이하 자연수 중 $p$의 배수가 아닌 것을 곱한 값을 $p^e$로 나눈 나머지.
  • $[1, p^e-1]$ 중 $p$의 배수가 아닌 것부터 $[p^e(k'-1)+1, p^ek'-1]$ 중 $p$의 배수가 아닌 것이 먼저 있을 것이다.
  • 그 후 $[p^ek'+1, p^ek'+r']$ 중 $p$의 배수가 아닌 것이 여기에 있다. 이를 모두 종합하면?
  • $val[p^e-1]^{k'} \cdot val[r'] \pmod{p^e}$이라는 값이 튀어나온다. 결국 기존 주제 3과 거의 다를 바가 없다. 
  • 또한, $val[p^e-1]$이 $\pmod{p^e}$로 $1$ 또는 $-1$이라는 결과 역시 알려져 있다. 그러니 $k'$의 홀짝성만 봐도 ok.

이 알고리즘은 $\mathcal{O}(p^e)$ 급의 계산이 필요하다. 하지만 실제로는 $\mathcal{O}(pe)$ 급으로 계산이 가능하다. 매우 어렵다.

이 강화된 알고리즘을 공부하고 싶다면, HYEA님의 이 글과 원본이라고 볼 수 있는 min_25의 이 글을 보자.


이항계수

여전히 $k \le n/2$만 보는 아이디어는 유효하다.


Case 1. $k<p$이며 $k$가 작은 경우. 

이 경우, $1, 2, \cdots , k$ 각각의 $\pmod{p^e}$에 대한 modular inverse가 존재하므로 아이디어가 그대로 적용된다.


Case 2. 작은 $n, k$에 대하여 전부 전처리를 하고 싶은 경우.

이 경우, 파스칼의 삼각형을 이용하여 DP와 같은 방식으로 계산을 할 수 있다. 정수론이 필요없는 영역이다.


Case 3. $p$가 작은 경우, 즉 $\mathcal{O}(p^e)$ 또는 $\mathcal{O}(pe)$ 알고리즘이 동작할 수 있는 경우.

이 문제를 Lucas 정리로 풀려고 하면 (즉, Lucas 정리의 확장판을 찾으려 한다면) 상당히 고생할 것이다. 적어도 내 기억은 즐겁지 않았다.

하지만 주제 3을 이용한 풀이를 생각한다면, 기존 Case 3과 완전히 동일한 풀이가 적용된다. 주제 3을 여기에 수록한 이유다.

 


이제 $\pmod{p^e}$에 대한 이야기를 끝냈으니, 이론적으로 모든 내용을 다 했다. 하지만 몇 가지 더 이야기거리가 남았다.

  • 중국인의 나머지 정리는 항상 강조하지만 수학에서나 정수론 알고리즘에서나 매우 중요한 결과고 쓸모도 많다.
  • 하지만 이걸 제대로 써먹으려면 (즉, prime power로 문제를 쪼갠 후 다시 합치는 방식을 적용) 소인수분해를 해야 한다.
  • 기본적으로 소인수분해는 매우 어려운 문제다. CP/PS 환경에서는 보통 커봐야 $2^{64}$인 수를 다루니 Pollard-Rho 알고리즘을 배우면 조금 낫다.
  • 하지만 Pollard-Rho가 엄청 빠른 알고리즘인 것도 아니다. 즉, 일반적인 $\pmod{M}$에 대해서 사용할 수 있는 알고리즘을 더 다룰 필요가 있다.
  • 당연한 것이지만 아래에서 다룰 "일반적인 알고리즘"은 위에서 다룬 $\pmod{p}$, $\pmod{p^e}$의 경우에서도 사용할 수 있다. 
  • 아래에서 다룰 알고리즘들은 modulus가 갖는 소인수 $p$나 prime power $p^e$의 크기 등에 영향을 받지 않는다.

팩토리얼


주제 1. $\mathcal{O}(n)$ 알고리즘은 당연히 아직도 유효하다.

주제 2. 팩토리얼만 구하는 $\mathcal{O}(n)$ 알고리즘은 당연히 아직도 유효하다.

modular inverse를 전부 구하는 문제가 잘 정의되려면 $1, 2, \cdots , n$이 전부 $M$과 서로소여야 한다.

이 경우, 방법 1, 3이 잘 적용된다. (소수, 소수의 거듭제곱 여부를 사용한 적이 없다)

주제 3. 일반적인 modulus $M$에 대해서 논의하기 애매한 주제다. 이 때문에 이항계수를 계산하기가 어렵다.


이항계수

Case 1. $k$가 작고, $1, 2, \cdots , k$ 전부가 $M$과 서로소인 경우.

이 경우에는 기존 알고리즘을 그대로 적용할 수 있다. $\pmod{M}$에 대한 modular inverse를 취할 수 있기 때문이다.

Case 2. 파스칼의 삼각형을 그대로 적용할 수 있다. 이미 말했지만, 이 부분은 정수론이 아니라 DP에 가깝다.

Case 3. 하기 힘들다. 애초에 각 소수가 작다는 가정이 전제되는 만큼, 소인수분해가 필요하다.

Case 4. 특수한 경우를 하나 더 소개한다. 바로 $n, k$가 둘 다 적당히 작은 경우다.

이 경우, 앞서 소개한 에라토스테네스의 체를 이용한 로그 시간 소인수분해를 하여 문제를 해결할 수 있다.

즉, $1, 2, \cdots , n$ 전부를 소인수분해 할 수 있으니, $\displaystyle \binom{n}{k} = \frac{n \cdot (n-1) \cdots \cdot ( n-k+1)}{1  \cdot 2 \cdots \cdot k}$ 역시 그 분모/분자를 전부 소인수분해 할 수 있다.

이를 통해서 아예 $\displaystyle \binom{n}{k}$의 소인수분해를 얻을 수 있다. 이제 나눗셈에 대한 걱정이 아예 없으니, 단순 계산이 가능해진다. 끝.