I participated in 0x41 CTF as a member of Super Guesser, and we reached 2nd place.

 

This is a very rushed writeup, hope to patch it up after I take a sleep...

 

 

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
 
def key_expansion(self, key):
        keys = [None* 5
        keys[0= key[0:4+ key[8:12]
        keys[1= key[4:8+ key[12:16]
        keys[2= key[0:4+ key[8:12]
        keys[3= key[4:8+ key[12:16]
        keys[4= key[0:4+ key[8:12]
        return keys
 
def apply_sbox(self, pt):
    ct = b''
    for byte in pt:
        ct += bytes([sbox[byte]])
    return ct
 
def apply_perm(self, pt):
    pt = bin(int.from_bytes(pt, 'big'))[2:].zfill(64)
    ct = [None* 64
    for i, c in enumerate(pt):
        ct[perm[i]] = c
    return bytes([int(''.join(ct[i : i + 8]), 2for i in range(0len(ct), 8)])
 
def apply_key(self, pt, key):
    ct = b''
    for a, b in zip(pt, key):
        ct += bytes([a ^ b])
    return ct
 
def handle(self):
    keys = self.key_expansion(key)
    for i in range(65536):
        pt = os.urandom(8)
        ct = pt
        ct = self.apply_key(ct, keys[0])
        for i in range(ROUNDS):
            ct = self.apply_sbox(ct)
            ct = self.apply_perm(ct)
            ct = self.apply_key(ct, keys[i+1])
        self.send(str((int.from_bytes(pt, 'big'), int.from_bytes(ct, 'big'))))
cs

 

Summary of the Challenge

Basically, we are given a 4-round SPN cipher, and we have to completely find the key.

To do this, we are given $2^{16}$ known plaintext/ciphertext pairs. We also know the sbox/pbox.

 

The Attack Idea, and the SBOX

Since this is a known plaintext/ciphertext attack, we can think about linear cryptanalysis.

To do that, we have to exploit the SBOX's linear properties. After some researching on past linear cryptanalysis challenges in CTFs, I ended up at this writeup of zer0SPN by NGG. I ran the first code block to find that the SBOX was horrible beyond belief. Denoting $\text{bit}(n, k)$ for $0 \le n < 2^8$ as the $k$th bit from the left when we write $n$ as a $8$-digit binary integer, we can see that $\text{bit}(n, k)$ and $\text{bit}(sbox[n], k)$ are heavily correlated. This is a very simple output compared to the zer0SPN challenge, where several bits come into play at once.

This is a very promising outcome, and shows we are in the right direction. I actually thought I was done at this point :P

 

Attack 1 : Recovering Linear Equations of the Key Bits

Now what we can do is view the SBOX as a simple XOR. 

 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def BIT(v, k):
    cc = bin(v)[2:].zfill(8)
    return ord(cc[k]) - ord('0')
 
invert = []
 
for i in range(08):
    cnt_0, cnt_1 = 00
    for j in range(0256):
        if BIT(j, i) == BIT(sbox[j], i):
            cnt_0 += 1
        else:
            cnt_1 += 1
    if cnt_0 > cnt_1:
        invert.append(0)
    else:
        invert.append(1)
cs

 

Basically, with the above code, we can regard the SBOX as

  • if $invert[i] = 0$, the SBOX does nothing to the $i$th bit
  • if $invert[i]=1$, the SBOX flips the $i$th bit

If we use this idea, then literally everything we do is just permuting and XOR'ing each bit.

Of course, the SBOX is not "perfectly" an XOR, but it will be biased towards it due to the properties we found.

This gives us a method to find 64 linear equations of the key bits in $\mathbb{F}_2$. Here's how we do it.

 

Let's think about the very first (0th) bit of the plaintext. This bit will

  • get XOR'ed by the first key bit of the first key
  • SBOX'ed (which can be approximated by XOR'ing $invert[0]$)
  • PBOX'ed (moved to a new position, where we know)
  • and repeated like that for four times

note that we know the values we "XORed" in the SBOX, and we know the location of the initial 0th bit.

If the $0$th bit ended up in the $u$th bit after the four PBOX applications, we can see that

  • plaintext's $0$th bit gets XORed to some bits we know and some key bits we don't know and becomes the ciphertext's $u$th bit

Of course, this is all an approximation. The fact that remains true, is that the value of $ptxt[0] \oplus ctxt[u]$ will be biased towards either $0$ or $1$. From this, we can retrieve the XOR of key bits that was applied to the bit that was originally the 0th bit of the plaintext.

Doing the same thing for each of the 64 bits, we get 64 equations. Here's the code for the first part.

 

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
def whi(idx, i): # what key bit am I actually using?
    if idx % 2 == 0:
        if i < 32:
            return i
        else:
            return 32 + i
    else:
        if i < 32:
            return 32 + i
        else:
            return 64 + i
 
for i in range(064):
    loc, add = i, 0
    myenc = []
    myenc.append(whi(0, loc)) # key XOR
    for j in range(15):
        add += invert[loc % 8# sbox
        loc = perm[loc] # pbox
        myenc.append(whi(j % 2, loc)) # key XOR
    myenc.append(add % 2)
    myenc.append(loc)
    arr.append(myenc)
 
= remote('185.172.165.118'4004)
pt, ct = [], []
 
print("[+] Receving Plaintext")
 
for i in range(065536):
    tt = r.recvline()
    tt = tt.split(b" ")
    A = bin(int(tt[0][1:-1].decode()))[2:].zfill(64)
    B = bin(int(tt[1][:-2].decode()))[2:].zfill(64)
    pt.append(A)
    ct.append(B)
 
ZERO, ONE = [], []
 
for i in range(064):
    fin = arr[i][-1# final location
    cnt_0, cnt_1 = 00
    for j in range(065536):
        st = ord(pt[j][i]) - ord('0')
        en = ord(ct[j][fin]) - ord('0')
        if st == (en + arr[i][-2]) % 2# XOR of the key bits is 0
            cnt_0 += 1
        else# XOR of the key bits is 1
            cnt_1 += 1
    print(cnt_0, cnt_1) # check bias
    if cnt_0 > cnt_1:
        ZERO.append(arr[i][:-2]) # sum of these = 0
    else:
        ONE.append(arr[i][:-2]) # sum of these = 1
 
cs

 

Attack 2 : Recovering the First Key XOR'ed

Here, we get the entire first key, giving us 64 more bits of information. This solves the challenge.

We will use the 0th bit of the plaintext as an example again. We will denote the final positions of the initial $i$th bit as $f[i]$. In other words, $f[i]$ is simply the P-box iterated four times. We already know that $ptxt[i]$ and $ctxt[f[i]]$ are correlated.

 

However, if we knew the first byte of the first key getting XOR'ed, we can get even better correlation.

In this case, we know how the first byte of the plaintext is going to behave for the first round of encryption - 

  • we know the key being XOR'ed 
  • we can now directly apply the SBOX without approximating
  • we can directly apply the PBOX without any issue

At this point, the "new" value of the initial $i$th bit will have a better correlation with $ctxt[f[i]]$. 

You can get this intuitively, since we have used "less" approximations, or you may use the piling-up lemma.

 

Therefore, we can now brute force the first byte, do the above computation, and see which byte gives the best correlation.

This can be done for each byte separately, giving our 8 key bytes, or 64 bits. Here's the code.

 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
for i in range(08): # ith byte
    print("[+] Guessing key", i)
    ideals = []
    for j in tqdm(range(0256)): # bruteforce
        cnt_ideal = 0
        for idx in range(08):
            cnt_0, cnt_1 = 00
            for whi in range(065536): # over ptxt/ctxt pairs
                fin_loc = arr[8 * i + idx][-1]
                addv = arr[8 * i + idx][-2- invert[idx]
                bt = BIT(sbox[int(pt[whi][8 * i : 8 * i + 8], 2) ^ j], idx) # the first round
                res = ord(ct[whi][fin_loc]) - ord('0')
                if bt == res:
                    cnt_0 += 1
                else:
                    cnt_1 += 1
            cnt_ideal += max(cnt_0, cnt_1) # the correlation
        ideals.append(cnt_ideal)
    mx = 0
    for j in range(0256): # max correlation
        mx = max(mx, ideals[j])
    print(ideals.index(mx))
cs

 

Finishing Touches

We can embedd every information we gathered as a system of linear equations over $\mathbb{F}_2$.

Solving this with SageMath, we see that there are $16$ solutions for the system. We try all and see which matches the ptxt/ctxt pair.

'CTF' 카테고리의 다른 글

UnionCTF Crypto Writeups  (2) 2021.02.21
CryptoHack All Solve, Again  (10) 2021.01.31
TetCTF 2021 Crypto Writeups  (1) 2021.01.03
PBCTF 2020 Crypto Writeups  (1) 2020.12.07
N1CTF 2020 Crypto Write-Ups  (0) 2020.10.18

This post attempts to (in a half-mathematical, half-intuitive way)

  • briefly explain zero-knowledge proofs and their definitions
  • show an example with Schnorr's protocol for proving knowledge of discrete log
  • introduce the integer commitment scheme of Damgard and Fujisaki (2001)
  • introduce diophantine sets, and diophantine proof of knowledge (2003)

I'm also attempting to write a demonstration for this paper. (with additional assumptions and simplifications) Here's the link.

 

1. Zero Knowledge Proofs

 

Here's a very well-written introduction for zero-knowledge proofs by Matthew Green

Zero-knowledge proofs are proofs that reveals nothing except for the assertion that the statement one is trying to prove is correct. For example, if I want to prove to a verifier that "I have this secret value $x$ satisfying some property", I want to do so without giving any extra information on $x$. At the same time, we want this proof to have the fundamentals of our intuition of a "proof" - that is, the verifier can be convinced of the proof if the statement is true, and the verifier can reject the proof if the statement is false.

 

To summarize, we want

  • Completeness : If statement is true, the verifier properly following the protocol can verify the proof.
  • Soundness : If statement is false, the verifier properly following the protocol can reject the proof with overwhelming probability.
  • Zero-knowledge : If statement is true, the verifier learns nothing except for the fact that the statement is true.

We call people (prover, verifier) that are properly following a protocol to be honest.

We call people (prover, verifier) that are not properly following a protocol to be malicious.

For example, a malicious prover may want to forge a "proof" for a statement that is not true.

Also, a malicious verifier may want to extract some extra knowledge from the proof by not following the protocol precisely.

 

So we have this "learns nothing except for the fact that the statement is true" thing : how do we formalize it in mathematical terms?

The proofs we will deal in this writeup will be interactive proofs - proofs where the prover and verifier interact with each other to see whether the prover's claim is true of not. Also, the proof protocol will not be deterministic. Therefore the "proof script" - the complete interaction of the honest prover and verifier - will follow some distribution. If we simulate this proof protocol "precisely" by only using the statement that is being asserted, this will intuitively be a good definition of "zero-knowledge". Summarizing, we have two distributions : 

  • The distribution of the actual "proof script" between the honest prover and the verifier
  • The distribution of the simulated "proof script" designed by only using the statement that is being asserted

 

Now we can truly define what it is to be "zero-knowledge" - 

  • Perfect Zero-Knowledge : No matter what verifier does, there's a simulated "proof script" that follows the exact same distribution with the actual "proof script" - so it is impossible to gather any extra knowledge, even with infinite computational resources
  • Statistical Zero-Knowledge : The same as perfect zero-knowledge, but now the two distributions does not need to be the same - only "close". How do we determine if two distributions are close? Check out statistical distance
  • Computational Zero-Knowledge : The same, but the two distributions are computationally hard to distinguish.
  • Honest-Verifier Zero-Knowledge : We now add the assumption that the verifier is honest. 
Dealing with malicious verifier is quite tricky, so we will deal with honest verifiers for the rest of this writeup. 
 
 
2. An example : Schnorr's Protocol
 
Here, we quickly describe Schnorr's Protocol for proving the knowledge of discrete logarithm.
We only consider the honest-verifier version, as it is much easier to understand the reasoning behind.
 
Consider a group $G$ with prime order $p$ - obviously it is a cyclic group. Prover & Verifier both know a generator $g \in G$ and some $y \in G$.
The discrete logarithm of $y$ over base $g$ is $x$ such that $g^x = y$, and it is generally hard to compute.

 

The prover now wants to prove their knowledge of $x$ such that $g^x = y$, without revealing any additional info on $x$.
To do this, Schnorr's protocol works as follows - all randoms are uniform random.
  • Prover generates a random $r \in [0, p)$ and sends $g^r$ to the verifier.
  • Verifier issues a random challenge $c \in [0, p)$ and sends $c$ to the prover.
  • Prover sends $z = (r + cx) \pmod{p}$ to the Verifier. 
  • Verifier checks if $g^z = g^r y^c$ is true. If not, the proof is false.
  • They repeat the process until the probability that a malicious prover got away with a fake proof is small as desired.
 
Now let us check the three properties we need under the honest-verifier assumption.
  • Completeness : If prover does know $x$, we see that $g^z = g^{r+cx} = g^r y^c$ is indeed true, so the verifier accepts it.
  • Soundness : Here's an intuitive explanation that can be outlined to a proof. If some prover can answer the verifier's challenge with a high probability, the prover should know the value of $(r + cx) \pmod{p}$ for many values of $c$. From these values, one can easily recover the value of $x$. Therefore, we have proved that the prover must actually know $x$ in order to answer the challenges well.
  • Zero-Knowledge : The actual proof script is $(g^r, c, (r + cx) \pmod{p})$ for i.i.d. $r, c \in [0, p)$. This is equal to the distribution of $(g^z y^{-c}, c, z)$ for i.i.d. $c, z \in [0, p)$. Therefore, we can perfectly simulate the proof script only using $g, y$ and the statement $g^x = y$. 

This proves that the Schnorr's Protocol is honest-verifier perfect zero knowledge. 

If the challenge is selected in $[0, 1]$ (binary challenge) the protocol can be proved to be perfect zero-knowledge.

 

 

3. Integer Commitment - Damgard & Fujisaki
 
Consider the following situation : you want to make a prediction about the future, but you don't want to reveal it immediately for whatever reason. What you want to do is wait until you know whether you got it correct, and if you did, you want to prove that you predicted the future correctly. How would one do this? One method (for one time use) is to use a cryptographic hash function.

 

Consider a (deterministic) hash function $H$ such that 
  • It is hard to generate a collision pair - i.e. $m \neq m'$ such that $H(m) = H(m')$.
  • It is hard to find a preimage - i.e. given $H(m)$, it is hard to recover $m$.
  • A small change in the message $m$ leads to a huge change in the hash $H(m)$.

Now what we can do is announce to the world that you have made a prediction with hash $H(m)$. 

People who see the value $H(m)$ cannot recover the value $m$ since it is hard to find a preimage. 
After that, you can claim that you predicted $m$, by showing that the hash of $m$ is indeed equal to $H(m)$.
Since it is hard to generate a collision pair, you can't claim that you predicted something different, like $m'$. 
 
This is a very simple scheme with some flaws (for example, you can't use it twice on the same message) but it's still cool since
  • Binding : Once you commit $m$, you cannot go back to another message $m'$.
  • Hiding : No one can recover the value of $m$ from $H(m)$, so it hides the value of $m$.

 

Commitment schemes are of similar nature. Here, we will specifically look at integer commitment.

We want the cool properties of binding and hiding, and we want them to be defined clearly. Here's how we do it.

  • Consider a deterministic commitment function $C$ that takes two integers as input.
  • To commit an integer $x$, we generate a random integer $r$ and use $C(x, r)$ as our commitment value.
  • The constraints on $x$ and the interval of $r$ to be selected from uniformly will be determined later.
  • Binding : It is difficult to find $x, x', r, r'$ such that $x \neq x'$ and $C(x, r) = C(x', r')$.
  • Hiding : For any $x, x'$, the distributions of $C(x, \cdot)$ and $C(x', \cdot)$ are equal/statistically close.
 
Now, what if we want even more out of our integer commitment scheme? Here's an additional property that would be cool if we have it.
  • Summation Protocol : If prover has three integers $a, b, c$ such that $a + b = c$ and their respective commitment $A, B, C$, they can prove that $a + b = c$ in a zero-knowledge way using only the commitment values $A, B, C$. 
  • Multiplication Protocol : If prover has three integers $a, b, c$ such that $ab = c$ and their respective commitment $A, B, C$, they can prove that $ab = c$ in a zero-knowledge way using only the commitment values $A, B, C$. 

In other words, the prover can claim that "these commited values actually have this sum/product relation" in a zero-knowledge way.

 

It turns out that there is such a method, and it was devised by Damgard and Fujisaki in this paper

To make explaining things much easier, we will cover only a part of the results shown in the paper. To be exact,

  • We will only cover a specific case of the originally much, much more general construction of the scheme.
  • We will not discuss the details of the proof of why the schemes are valid, but discuss some intuition & insights.
  • Some protocols outlined below are not from the paper - these are mostly my thoughts & guesses.
  • We will only consider the honest-verifier case - the transformation from honest-verifier case to the general case is dealt in the paper.
  • The paper for Diophantine Argument of Knowledge also deals with extensions of this scheme. We do not deal with them here.

 

Now we discuss the actual scheme by Damgard and Fujisaki. We first need to describe the setup process.

  • Setup - Part 1. The verifier takes two distinct, large safe primes $P, Q$ and let $n = PQ$. From the definition of safe primes, we can write $P = 2p+1$, $Q = 2q+1$ with $p, q$ being a prime. The verifier then constructs an element $h \in \mathbb{Z}_n^{\times}$ such that $h$ has order $pq$. From now on, we will work in the cyclic group $\langle h \rangle \le \mathbb{Z}_n^{\times}$ for our commitments. 
  • Setup - Part 2. Assume that $n$ is a $k$-bit modulus - then anyone who knows $n$ also knows that the cyclic group generated by $h$ has order less than $2^B$ for $B = k$. This is an important piece of information for statistical zero-knowledge.
  • Setup - Part 3. The verifier generates $\alpha$ randomly from $[0, 2^{B+k})$, and sets $g = h^{\alpha}$. The verifier sends $n, g, h$ to the prover, and the verifier proves the knowledge of $\alpha$ such that $g = h^{\alpha}$ in a zero-knowledge way. This can be done with a binary challenge process similar to Schnorr's protocol. It's important to note that since factorizing $n$ is hard, the prover does not know the order of $h$.
  • Note that the prover also knows the values of $B$ and $k$, since they can be easily derived from $n$.
  • Setup - Part 4. We also need an additional parameter $C$ - an higher value of $C$ will result in lower error probability, but at the same time we need $C$ to satisfy one of our assumptions below. Because we are dealing with a specific case of a much general construction in the paper, setting $C$ as high as $pq$ is not an issue. For more insight on selecting $C$, it is recommended to refer to the original paper.

 

Now that we know what we are working with - integers modulo $n$ as commitment, we can outline some assumptions.

While there are actually more assumptions involved, using the safe prime construction only leaves us with two - 

  • Strong Root Assumption : Basically, if one is given a random $y \in \langle h \rangle$, then it is hard to find $x \in \langle h \rangle$ and $t \ge 2$ such that $y = x^t$.
  • Small Order Assumption : Basically, it is hard to find a $\sigma$ and $b \in \langle h \rangle$ such that $0 < \sigma < C$ and $b \neq 1$ and $b^\sigma = 1$. 

In the safe prime construction, we can just replace the two with the Strong RSA assumption, alongside with a suitable value of $C$.

 

 

Now we discuss how to commit integers. To commit an integer $x$, the prover selects some $r \in [0, 2^{B+k})$ and makes a commitment as $C(x, r) = g^x h^r$. It can be proved that this commitment scheme has the (statistical) hiding property and the binding property. 

 

Also, from now on, we assume that the integers being commited are in some finite interval $[-T, T]$. This is a required assumption for proofs to be zero-knowledge. Of course, a malicious prover can try to commit integers that are not in this interval - therefore, the verifier must also ask for a proof that the integer is indeed in this interval. The method for this "range proof" will be outlined in the next section.

 

 

Now we discuss the proofs between the prover and the verifier. There are many things we want to prove, these include

  • Proving you know how to open : Assume the prover has some integer $x$ and its commitment $C$. The prover now wants to prove the knowledge of $x, r$ such that $C = C(x, r)$ - can it be done without leaking information on $x, r$?
  • Summation Protocol : If prover has three integers $a, b, c$ such that $a + b = c$ and their respective commitment $A, B, C$, they can prove that $a + b = c$ in a zero-knowledge way using only the commitment values $A, B, C$. 
  • Multiplication Protocol : If prover has three integers $a, b, c$ such that $ab = c$ and their respective commitment $A, B, C$, they can prove that $ab = c$ in a zero-knowledge way using only the commitment values $A, B, C$. 

 

Proving you know how to open

  • Keynotes : We imitate Schnorr's protocol - but the range we choose our parameters from needs to be selected carefully.
  • Step 0. Prover has some commitment $C = C(x, r)$ and would like to prove they know how to open.
  • Step 1. Prover chooses some random $y \in [0, TC 2^k)$ and $s \in [C 2^{B+2k})$, sends $d = g^y h^s$.
  • Step 2. The verifier issues a challenge $e \in [0, C)$ and sends it to the prover.
  • Step 3. Prover sends $u = y + ex$ and $v = s + er$ - verifier checks $g^u h^v = dc^e$

 

Summation Protocol

  • Keynotes : The key is that if $A = g^a h^{r_1}$, $B = g^b h^{r_2}$, $C = g^c h^{r_3}$ with $a + b = c$, we know the discrete logarithm of $C(AB)^{-1} = h^{r_3-r_1-r_2}$ over base $h$. Since the prover cannot calculate the discrete logarithm of $g$ over $h$, this can be a proof.
  • Step 0. Prover has three commitments $A, B, C$ of $a, b, c$ and would like to prove $a + b = c$.
  • Step 1. Prover first proves that they can open the commitments $A, B, C$ in a zero-knowledge way.
  • Step 2. Prover then proves the knowledge of discrete logarithm of $C(AB)^{-1}$ over base $h$ in a zero-knowledge way. Since we have assumed honest-verifier setting, it can be done in a way similar to Schnorr's protocol. Of course, we have to be careful about the range we choose our parameters from. The same applies for the Setup stage, where the verifier proves the knowledge of $\alpha$ to the prover. For details, refer to the original paper, page 5. This completes the proof of $a + b = c$, in a zero-knowledge way.

 

Multiplication Protocol

  • Keynotes : The key is that if $A = g^a h^{r_1}$, $B = g^b h^{r_2}$, $C = g^c h^{r_3}$ with $ab = c$, we have $C = A^b h^{r_3 - br_1}$.
  • Step 0. Prover has three commitments $A, B, C$ of $a, b, c$ and would like to prove $ab = c$.
  • Step 1. Prover chooses random $y \in [0, TC2^k)$, $s_2 \in [0, C 2^{B+2k})$, $s_3 \in [0, TC2^{B+2k})$.
  • Step 2. Prover sends $d_2 = g^y h^{s_2}$, $d_3 =  A^y h^{s_3}$ to the verifier.
  • Step 3. Verifier issues challenge $e \in [0, C)$ and sends it to the prover.
  • Step 4. Prover sends $u = y + eb$, $v_2 = s_2 + er_2$, $v_3 = s_3 + e(r_3 - br_1)$ to the verifier.
  • Step 5. Verifier checks if $g^u h^{v_2} = d_2 B^e$, $A^u h^{v_3} = d_3 C^e$ holds.

 

We will now (very) briefly sketch how the each protocols are proved to be zero-knowledge. The stuff going on is basically

  • Proving Completeness : Trivial - just substitute all values in and check if the equation is true.
  • Proving Soundness : "Standard proof techniques for soundness" + Assumptions (Strong Root, Small Order, Strong RSA, etc)
  • Proving Zero-Knowledge : Done by simulation - the choice of the intervals in which we select our parameters comes into play here.
  • Additional Notes : It's very important that the prover does not know the order of $h$ - for example, if the prover knew the order $pq$, they can forge a commitment of $x$ into a commitment of $x + pq$ easily. However, this condition makes it hard for the prover to truly generate a "uniform random distribution" (for example, in Schnorr's protocol, the prover generates the integer in $[0, p)$, where $p$ is the group order) - and this is why we need statistical zero-knowledge instead of the "perfect" variant. It's quite tricky...

 

 

4. Diophantine Sets and Diophantine Argument of Knowledge

 

Now let's think about what we can do with our integer commitment scheme. Working with integers in $[-T, T]$, we can 

  • Prove $a + b = c$ with only commitments, in a zero-knowledge way. 
  • Prove $ab = c$ with only commitments, in a zero-knowledge way.

It's clear that by using addition and multiplication protocols we can actually prove that we have computed an integer coefficient multivariable polynomial on some input. In other words, we can prove that $y = f(x_1, x_2, \cdots , x_n)$, for some integer coefficient polynomial that both prover & verifier know using only the commitment values for $x_i$'s and $y$ in a zero-knowledge way. Here's an easy example.

  • Prover has some values $x_1, x_2, x_3, x_4$ that satisfies $x_4 = x_1x_2 + x_3$. How to prove it?
  • Prover calculates $x_5 = x_1x_2$ and commits $x_1, x_2, x_3, x_4, x_5$.
  • Prover can now prove $x_5 = x_1x_2$ and $x_4 = x_3 + x_5$ in a zero-knowledge proof, only using commitments.

 

While that's cool and already has quite a lot of applications, it's quite limited. There are plenty of prepositions that one would like to prove that does not have the form of "I evaluated a integer coefficient polynomial"... - or are there?

 

To dig deeper here, we need the notion of Diophantine Sets. 

  • Diophantine Sets : A set $S \subset \mathbb{Z}^k$ is called Diophantine if and only if there exists a integer coefficient multivariate polynomial $P$ such that $v \in S \iff \exists w \in \mathbb{Z}^{k'} \text{  s.t.  } P(v, w) = 0$, i.e. if $v \in S$ iff $v$ is a part of a solution for a fixed Diophantine equation.
  • In the above definition, we call $P$ as a representing polynomial of $S$, and $w$ as a witness.

This notion leads us to a much more versatile approaches for a proof. If $S$ is Diophantine and I want to prove $v \in S$, I can

  • Find a representing polynomial $P$ of $S$ and publicize it along with the full formula & proof for $P$.
  • To prove $v \in S$, find a witness $w$. Prove $P(v, w) = 0$ using integer commitment scheme.

This is the fundamental idea of Diophantine Argument of Knowledge, abbreviated as DARK.

The obvious question is "what sets are Diophantine?" : after all, it should cover a lot of different sets for this to have practical meaning.

 

Here's one of the more surprising results in mathematics - 

  • Computably Enumerable Sets : A set $S$ is computably enumerable if and only if there is an algorithm that halts if the input is a member of $S$, and runs forever if otherwise. Equivalently, a set $S$ is computably enumerable if there is an algorithm (not necessarily halts) that enumerates the members of $S$. Obviously, a lot of sets we deal with are computably enumerable.
  • Matiyasevich's Theorem (MRDP Theorem) : [Diophantine Sets] = [Computably Enumerable Sets]

so a metric ton of sets are Diophantine. However, to put into practical use, more questions must be asked.

  • How do we even find the suitable representing polynomial $P$? How do we prove it?
  • If we found $P$, how do we find the witness $w$? Is the size (number of bits) of $w$ reasonable/practical?

Therefore, we deal with special Diophantine sets that 

  • have representing polynomial $P$ (we can actually construct)
  • has a polynomial time algorithm to compute the witness $w$, which has polynomial length.

In this paper, we narrow down our choices even further, by adding the constraint

  • the length of our witness has to be subquadratic to the length of the input

which is important for practical uses. Now our question is "what can we prove under these constraints?".

 

 

5. Tricks, Range Proofs, Exponential Relations, and Bounded Arithmetic

 

Tricks

Here we outline some tips and tricks that we will use later on. Suppose we have two Diophantine sets $S_1$ and $S_2$, with respective representing polynomials $P_1, P_2$. How do we deal with $S_1 \cap S_2$ or $S_1 \cup S_2$? These will be our way to deal with and/or operations on prepositions. 

  • $S_1 \cap S_2$ : We use $P_{\cap}(v, w_1, w_2) = P_1(v, w_1)^2 + P_2(v, w_2)^2$ as our polynomial.
  • $S_1 \cup S_2$ : We use $P_{\cup}(v, w_1, w_2) = P_1(v, w_1) \cdot P_2(v, w_2)$ as our polynomial.

It's relatively clear that these two work as desired, and witnesses can be transported over as well.

 

Range Proofs

We have a homework from chapter 3 - we have to prove a value is in some range for the scheme to work.

To prove $-T \le x \le T$, it suffices to prove $T-x \ge 0$ and $T+x \ge 0$. Therefore, it suffices to have a scheme that proves an integer is nonnegative. To do this, it suffices to have a nice representing polynomial $P$ for nonnegative integers, with witness easy to find.

  • Construction : We use the polynomial $P(x, a, b, c, d) = x - a^2 - b^2 - c^2 - d^2$.
  • Proof (representation) : Obviously, if $P(x, a, b, c, d) = 0$, we have $x = a^2 + b^2 + c^2 + d^2 \ge 0$. On the other hand, if we have $x \ge 0$, we can write it as $a^2 + b^2 + c^2 + d^2$ by Lagrange's Four Square Theorem.
  • Proof (efficiency) : It is clear that the witness size is not an issue here. We show that it's feasible to find $a, b, c, d$ such that $a^2 + b^2 + c^2 + d^2 = x$. A quick outline of an (probabilistic) algorithm is as follows - 
  • Algorithm : If $x$ is small, we can simply brute force over $a, b, c, d \in [0, \sqrt{x}]$. If $x$ is large, take $c$ as a random value in $[0, \sqrt{x}]$ and $d$ as a random value in $[0, \sqrt{x-c^2}]$. Now, we pray that $x - c^2 - d^2$ is actually a prime of the form $4k+1$. Heuristically, with Prime Number Theorem, we can see that there exists a good probability that this holds. Now all we need to do is write $x-c^2-d^2$ as a sum of two squares. There exists a few fast algorithms that writes a $4k+1$ prime as a sum of two squares - check out this cool thread.

Now we can argue stuff like $a \ge b$ in our proofs, and it gives us much more versatility! 

 

Exponential Relations

Here, our goal is to prove $c =  a^b$. Our views on representation polynomial are a bit different since we now have a view that is focused on practical applications - we will change it as follows. This turns out to be important for proving $c = a^b$ zero-knowledge.

  • If $v \in S$, we should be able to efficiently find a polynomial length $w$ such that $P(v, w) = 0$.
  • If $v \notin S$, we should not be able to efficiently find a polynomial length $w$ such that $P(v, w) = 0$ - one way to prove this "not able to efficiently find" thing, is to prove that all the fake "witness" $w$ with $P(v, w) = 0$ has very large length.

We can now enjoy the following theorem, and I will not even try to type this out or code this.

Let's decipher this. Since the expressions (E1-E10) are all prepositions that can be written as a polynomial equation, we can combine them as a one big integer coefficient multivariate polynomial $P$ with our "trick". Basically, if $\mu_1^{\mu_2} = \mu_3$, we can efficiently find our witness for $P$ that has subquadratic length. If $\mu_1^{\mu_2} \neq \mu_3$, either the witness for $P$ does not exist, or the witness has a large length, $\Omega(\mu_3 \log \mu_3)$ - which makes it impractical to use as a fake "proof" that $\mu_1^{\mu_2} = \mu_3$. We can also take care of edge cases like $\mu_1=0, 1$, $\mu_2 = 0, 1, 2$, $\mu_3=0$ easily with our "trick" as well. This shows that we can completely prove $a^b = c$ with DARK, zero-knowledge. 

 

Bounded Arithmetic

Now that we have a lot of tools at hand, we can take a look at what we can do.

  • Bounded Arithmetic : We work over the nonnegative integers, and only use the following operations
  • $0$, $\sigma$, $+$, $\cdot$, $\le$, $-$, $|x|$, $\lfloor x/2 \rfloor$, $\lfloor x/2^k \rfloor$, $\sharp$
  • $0, +, \cdot , \le , \lfloor x/2 \rfloor , \lfloor x/2^k \rfloor$ have usual definitions.
  • $\sigma(x)$ is simply $x$'s next integer, $x + 1$. 
  • $-$ is a bit different because we are working over nonnegative integers, $x-y$ is actually $\max(x-y, 0)$.
  • $|x|$ is the bitwise length of $x$, i.e. $\lfloor \log_2 (x) \rfloor + 1$, excluding the case $x = 0$. $x \sharp y$ is defined as $2^{|x| \cdot |y|}$

Since we now have and/or of prepositions, addition, multiplication, exponentiation ready to be proved in a zero-knowledge way, we can combine these with relative ease to show that all of those operations can be used, and proved in a zero-knowledge way. It seems that bounded arithmetic is quite researched, and a lot of useful prepositions can be written using the bounded arithmetic language.

 

Here's an outline of proof - 

  • $0, +, \cdot$ : We already know how to prove these by our integer commitment scheme.
  • $ \le $ : We studied how to do range proofs - simply proof $b - a \ge 0$ to show $a \le b$. 
  • $ - $ : $a - b = c$ is either ($a = b+ c$ and $c \ge 0$) or ($a \le b$ and $c = 0$), so combine.
  • $|x|$ : Deal with the case $x=0$ separately, and write $y = |x|$ as $2^{y-1} \le x \le 2^y-1$. Now do range proofs & exponentiation.
  • $\lfloor x/2 \rfloor$ : $y = \lfloor x/2 \rfloor$ is simply $x = 2y$ or $x = 2y + 1$, so combine them.
  • $\lfloor x/2^k \rfloor$ : $y = \lfloor x/2^k \rfloor$ is simply $x = 2^ky + z$ and $0 \le z \le 2^k-1$, so combine them.
  • $\sharp$ : Since we can do proofs for $|x|$ and proofs for exponentiation, combining is quite trivial.

There are definitely more content in this paper, including some practical applications and reduction of proof size.

There's also a relatively new paper in 2019 that is much harder and very interesting. Might try to read this one...

11단원의 내용으로 내가 알고 있는 PS/CP 정수론 분야 지식은 (특히, 문제풀이에 직접적인 내용은) 다 정리한 것 같다.

처음에는 그래도 친절하게 설명하려고 노력했는데, 후반부에 가니까 슬슬 다른 일의 우선순위가 높아져서 급하게 작성한 것 같아 아쉽다.

  • 초반 부분 : 이미 자료가 있으니까 간략하게 써야지 ㅋㅋㅋ
  • 후반 부분 : 내가 다른 일이 있으니까 간략하게 써야지 ㅋㅋㅋ

가 된 것 같아서 조금 그렇긴 한데 ㅋㅋ; 그래도 시간투자를 충분히 하면 이해할 수 있을 정도로는 쓴 것 같다. 나중에 코드, 연습문제도 올릴거고..


일단 이 내용을 다 이해했다면 (또는 관심있는 부분까지 다 이해했다면) 할 수 있는 것은 대강

  • 백준 문제풀기. 문제집 제작을 내가 하면 더 제대로 할 수 있을 것이고, 아니어도 solved.ac나 이 블로그 수학 PS일지를 보면서 문제를 고를 수 있다.
  • 진짜 고수들이 쓴 자료 읽기. 여러 의미를 담고 있는 말인데, 정리하자면
  • 1. 정수론 교과서 읽기. 이유를 설명하지 않고 넘긴게 많으니, 이를 이해하고 싶다면 책을 사서 읽는 것이 제일 빠르다.
  • 2. PS/CP 정수론 진짜 고수가 쓴 자료 읽기. 예로, kipa00님의 NTA, HYEA님의 소멤 글, min_25의 블로그, adamant의 블로그
  • adamant의 블로그는 PS 수학 전반에 대한 다양한 내용이 있다. 특히 FFT 및 다항식 처리 관련 자료가 어마어마하다.
  • HYEA님의 글과 min_25의 블로그에는 더 난이도 높은 알고리즘들이 많이 나와있다. 또한, HYEA님의 글 중 min_25의 글 번역본도 있다.
  • kipa00님의 NTA는 수학적 깊이가 상당하고, computational number theory의 여러 부분을 다룬다. 수학하는 사람이면 읽는 게 좋을듯.
  • 혹시 NTA를 다 읽고 컨텐츠가 부족하다고 느낀다면 (ㅋㅋ;) Victor Shoup의 computational number theory 책을 읽자.
  • 3. Project Euler 풀기. 여기에는 고수들이 가득하고, thread 등에서 풀이나 아이디어 공유가 많다. 
  • 특히, 10단원, 11단원의 내용을 확장/강화하는 토론들이 많이 이루어졌다. 대표적인 예시로 baihacker의 블로그를 참고하자. 
  • 개인적으로 중국/일본에 비해 한국에서 Project Euler를 푸는 사람이 적다는 것에 많은 아쉬움을 느끼고 있다.
  • 4. PS/CP scope 벗어나기. 이제 $2^{64}$ 미만의 수만 다루고, 짧은 시간 안에 답을 내야한다는 가정도 버리자.
  • 아예 Computational Number Theory를 공부할 수도 있고, (NTA 읽는 것도 방법) 그 대표적인 활용처인 암호학을 공부하는 것도 좋다.
  • 개인적으로는 PS/CP 정수론 및 PS/CP에서 겪은 문제풀이 경험이 CTF 암호학으로 넘어가는 것에 정말정말 많은 도움을 주었다.


시간이 없어서 글을 제대로 못 쓰기는 했는데 그래도 질문 받을 시간은 있으니 필요할 때 댓글 달아주세요 :)

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


이 글에서는 다음 내용을 다룬다.

  • multiplicative function의 정의, 성질
  • Dirichlet Convolution의 정의와 성질
  • xudyh's sieve, min_25 sieve를 이용한 multiplicative function의 prefix sum

이 내용을 다루는 튜토리얼 중 가장 접근 가능한 것은 역시 이 Codeforces 블로그. 나도 여기서 배웠다.


multiplicative function

  • 기존 정의를 생각하면, multiplicative function이란 $\text{gcd}(a, b)=1$일 때 $f(ab)=f(a)f(b)$인 함수 $f$를 말한다.
  • 특히, 우리가 다룰 함수들은 $f(1)=1$을 만족한다. ($f=0$ 같은 특수한 경우를 고려하지 않겠다는 것이다)
  • 이제 $n$을 소인수분해하고, $n=p_1^{e_1} \cdots p_k^{e_k}$라고 쓰자. 그러면 $f(n) = f(p_1^{e_1}) \cdots f(p_k^{e_k})$가 성립한다.
  • 이는 prime power에 대한 $f$ 값들만 결정하면 모든 자연수에 대한 $f$ 값이 자동 결정이라는 뜻이다.
  • 이 방법으로 생각하는 것이 아마 multiplicative function을 이해하는 가장 쉬운 방법일 것이다. 예시를 보자.
  • 오일러 파이 함수 : $f(p^k) = p^{k-1}(p-1)$으로 정의된 multiplicative function
  • mobius 함수 : $f(p)=-1$, $f(p^k)=  0$ for $k \ge 2$로 정의된 multiplicative function
  • 약수의 개수 함수 : $f(p^k) = k+1$로 정의된 multiplicative function
  • 약수의 합 함수 : $f(p^k) = 1 + p + \cdots + p^k$로 정의된 multiplicative function
  • identity 함수 : $f(p^k) = p^k$로 정의된 multiplicative function
  • $[n=1]$ 함수 : $f(p^k) = 0$으로 정의된 multiplicative function
  • $f=1$ 함수 : $f(p^k) = 1$으로 정의된 multiplicative function
등등, prime power만 보고 모든 값을 얻을 수 있는 함수들을 나열할 수 있을 것이다.

이제 multiplicative function의 여러 성질을 어렵지 않게 증명할 수 있다. 예시를 하나 보인다.

사실. $f(n)$이 multiplicative라면, $g(n) = \sum_{d|n} f(d)$ 역시 multiplicative.
증명. $g(p^k) = 1 + f(p) + \cdots + f(p^k)$인 multiplicative function이 됨을 보일 수 있다. 예를 들어, $$ g(p_1^{e_1} \cdots p_k^{e_k}) = \sum_{d_i \le e_i} f(p_1^{d_1} \cdots p_k^{d_k}) = \sum_{d_i \le e_i} f(p_1^{d_1}) \cdots f(p_k^{d_k}) = \sum_{d_1 \le e_1} f(p_1^{d_1}) \cdot \sum_{d_2 \le e_2} f(p_2^{d_2}) \cdots \sum_{d_k \le e_k} f(p_k^{d_k})$$이고 이는 정의상 $g(p_1^{e_1}) \cdots g(p_k^{e_k})$와 동일하다. 증명 끝. 이제 본격적인 주제로 들어가자.


Dirichlet Convolution

  • Dirichlet Convolution -> Dirichlet Series -> Analytic Number Theory의 연결 관계로 중요한 주제인 것으로 안다.
  • 하지만 PS/CP에서는 Dirichlet Convolution 그 자체가 중요한 경우는 거의 못 봤고, 아래에서 다룰 xudyh's sieve에서 등장한다.
  • 하지만, 실제로 Dirichlet Series와 같은 방식으로 PS/CP 문제풀이를 하는 경우를 Project Euler에서 봤다. 이건 나도 잘 모른다 ㅋㅋ
  • 정의. 두 함수 $f, g$에 대하여, 그 Dirichlet Convolution $(f * g)$를 $(f * g)(n) = \sum_{d|n} f(d)g(n/d)$로 정의한다.
  • 사실. $f, g$가 multiplicative 하다면, 그 Dirichlet Convolution $(f * g)$ 역시 multiplicative 하다.
  • 증명. 또 $(f * g)(p^k) = f(1)g(p^k) + f(p)g(p^{k-1}) + \cdots + f(p^k)g(1)$로 두고, 위에서 보인 예시처럼 하면 된다.
예시를 몇 개 들어보자.
  • $f(n) = n$, $g(n)=1$이면 $(f * g)(n)= \sigma(n) = \sum_{d|n} d$ - 약수의 합
  • $f(n)=1$, $g(n)=1$이면 $(f*g)(n) = \tau(n) = \sum_{d|n} 1$ - 약수의 개수
  • $f(n) = \mu(n)$, $g(n)=1$이면 $(f*g)(n) = \sum_{d|n} \mu(d) = [n=1]$ - mobius function
  • $f(n) = \phi(n)$, $g(n)=1$이면 $(f * g)(n) = \sum_{d|n} \phi(d) = n$
  • $f(n) = \mu(n)$, $g(n) = n$이면 $(f * g)(n) = \sum_{d|n} \mu(d) (n/d) = \phi(n)$
  • 특히 - mobius inversion을 제대로 표현하는 방법은 역시 $h = (f * 1)$이면 $f = (h * \mu)$라는 것이다.
  • $f(n) = n^k \phi(n)$, $g(n) = n^k$면 $(f * g)(n) = n^{2k+1}$ 등.

이제 메인 주제인, multiplicative function의 prefix sum을 계산하는 것으로 넘어가자.

xudyh's sieve
  • multiplicative function $f$가 있다. 우리의 목표는 $\sum_{i=1}^n f(i)$를 $\mathcal{O}(n^{2/3})$에 계산하는 것이다.
  • 적당한 multiplicative function $g$가 있어, $g$의 prefix sum과 $(f * g)$의 prefix sum이 계산하기 쉽다고 가정하자.
  • 함수 $f$에 대하여, $s_f(n) = \sum_{i=1}^n f(i)$라고 정의한다. 이제 준비 끝. 열심히 식을 조작해보자.

우선 Dirichlet Convolution의 정의를 사용하여 $$s_{f * g}(n) = \sum_{i=1}^n (f * g)(i) = \sum_{i=1}^n \sum_{d|i} f(i/d)g(d)$$를 얻는다. 이제 $d$에 대하여 식을 정리하면, 가능한 $i/d$의 값이 $1$부터 $\lfloor n/d \rfloor$이므로 $$ s_{f * g}(n) = \sum_{d=1}^n g(d) \sum_{k=1}^{\lfloor n/d \rfloor} f(k) = \sum_{d=1}^n g(d) s_f(\lfloor n/d \rfloor)$$이다. $g(1)=1$이니, 이제 이 식을 $s_f(n)$에 대하여 정리하면 결과적으로 $$s_f(n) =  s_{f * g}(n) - \sum_{d=2}^n g(d) s_f(\lfloor n/d \rfloor) $$

  • 계속해서 등장하지만, $\lfloor n/d \rfloor$ 형태를 갖는 서로 다른 자연수의 개수는 $\mathcal{O}(\sqrt{n})$개이다.
  • 또한, $\lfloor n/d \rfloor$의 특정 값에 대응되는 $d$의 값은 하나의 구간을 이룬다.
  • 특히, $g$의 prefix sum이 쉽게 계산될 수 있다고 가정했으므로, 그 구간에 대한 $g$ 값의 합 역시 쉽게 구할 수 있다.
  • 여기서는 당연히 $\sum_{i=l}^r g(i) = s_g(r) - s_g(l-1)$이라는 식을 사용한다. 
  • 그러니, 우리가 실제로 계산해야 하는 것은 $s_f(\lfloor n/d \rfloor)$ 값이다.

첫 방식은 $\lfloor n/d \rfloor$ 형태의 값들에 대해 $s_f$를 정직하게 DP로 계산하는 것이다.

$s_f(x)$를 계산하기 위해 드는 시간복잡도가 $\mathcal{O}(\sqrt{x})$이니, 우리에게 필요한 시간은 총 $$ \sum_{i=1}^{\sqrt{n}} \sqrt{i} + \sum_{i=1}^{\sqrt{n}} \sqrt{n/i} = \mathcal{O}(n^{3/4})$$임을 알 수 있다. 이제 지수를 또 $3/4$에서 $2/3$으로 줄여보자. 이를 위해서는 $f$가 multiplicative임을 이용한다.

  • 즉, 에라토스테네스의 체나 Linear Sieve를 사용하여, 작은 $i$들에 대한 $f(i)$를 전부 전처리하자.
  • 그러면 자연스럽게, 부분합을 계산하여 작은 $i$들에 대한 $s_f(i)$ 역시 전부 전처리할 수 있다.
  • 이번에는 $n^{2/3}$ 이하의 $i$에 대한 $s_f(i)$를 전처리하고, $s_f(n)$을 구하기 위한 점화식을 top-down으로 적용하자.
  • 이미 전처리된 값은 바로 반환하고, 메모이제이션을 이용하여 중복 계산을 막는다. 그러면 시간복잡도는 $$ \mathcal{O}(n^{2/3}) + \sum_{i=1}^{n^{1/3}} \sqrt{n/i} = \mathcal{O}(n^{2/3})$$이 된다. 이 방법이 xudyh's sieve이다. 이제 문제는 $s_f$를 구하고 싶을 때, $g$를 어떻게 잡냐는 것이다.

일단 $g$를 잡는 가장 좋은 방법은 위에서 선보인 예시들을 다 참고하는 것이다. 구글 검색하면 예시는 더 나온다.

특히, $1$, $n$, $n^k$ 등 $n$에 대한 다항식으로 나타나는 함수들은 그 prefix sum을 구하기가 상대적으로 쉽다는 점을 이용한다.

어려운 문제들을 풀기 시작하면, $g$를 잡기 위해서 전략적으로 각 prime power들에 대해 $g$의 값을 부여해야 할 수도 있다.

일단 위 문제의 가장 자주 보이는 예시는, $\phi(n)$의 합을 구하기 위하여 $g(n) = 1$, $(f * g) (n) =  n$을 이용하는 것이다.


구현 노트.

  • 구현 자체에 크게 어려운 점은 없으나, 역시 $\lfloor n/d \rfloor$가 매우 클 수 있으니 C++의 std::map을 이용해야 한다.
  • 하지만 std::map을 사용하면 로그가 하나 더 붙는다. 로그를 붙이지 않고 계산을 할 수는 없을까?
  • 애초에 $\lfloor n/d \rfloor \le n^{2/3}$이면 바로 값을 리턴하므로, $\lfloor n/d \rfloor > n^{2/3}$인 경우만 생각하자. 이 경우, $\lfloor n/d \rfloor$ 대신 그냥 $d$를 key 값으로 사용할 수 있다. 
  • 그러면 key 값이 $n^{1/3}$ 이하로 매우 작으니, std::map 대신에 그냥 배열을 사용해도 무방하다. 물론, 이 경우 초기화 문제에 주의해야 할 것이다.


min_25 sieve

  • 우선 필요한 정의부터 나열하고 시작하도록 하자. 모든 설명은 이 글을 따른다.
  • 이제부터 모든 나눗셈은 C++ 방식의 정수 나눗셈이다. 자연수만 다루니 부호는 신경쓰지 말자. (10장에서도 이럴 걸 ㅋㅋ)
  • $p_k$는 $k$번째 소수. 1-index를 사용한다. 즉, $p_1, p_2, p_3, \cdots = 2, 3, 5, \cdots$.
  • $lpf(n)$은 $n$이 갖는 최소의 소인수. 단, $n=1$인 경우 $lpf(1)=1$로 정의한다.
  • $F_{prime}(n)$은 $\sum_{2 \le p \le n} f(p)$이며, 이 합은 소수 $p$에 대한 것이다.
  • $F_k(n)$은 $\sum_{p_k \le lpf(i), 2 \le i \le n} f(i)$이다. 특히, $\sum_{i=1}^n f(i) = F_1(n) + f(1) = F_1(n)+1$이다.
  • 즉, $F_k(n)$은 최소 소인수가 $p_k$ 이상인 자연수에 대한 $f$ 값의 합이다. 이제 본격적으로 시작.

식 조작을 하자. 각 자연수의 최소 소인수가 무엇인가, 그리고 그 소인수를 몇 개 가지고 있는가로 분류하면, $$F_k(n) = \sum_{p_k \le lpf(i), 2 \le i \le n} f(i) = \sum_{k \le i, p_i \le n} \sum_{c \ge 1, p_i^c \le n} f(p_i^c) \left( 1 + F_{i+1}(n/p_i^c) \right) $$과 같다. 즉, $p_i^c$와 $1$의 곱, 또는 최소 소인수가 $p_{i+1}$이면서 $n/p_i^c$ 이하인 자연수의 곱으로 분류한다. 


한편, $n < p_k$면 $F_k(n) = 0$임을 사용하고, 적당히 식을 정리하면 $$ F_k(n) = \left( \sum_{k \le i, p_i^2 \le n} \sum_{c \ge 1, p_i^{c+1} \le n} \left( f(p_i^{c+1}) + f(p_i^c) F_{i+1}(n/p_i^c) \right) \right) + \sum_{k \le i, p_i \le n} f(p_i) $$를 얻는다. 여기서는 기존 식에서 $p_i^2 > n$이면 나오는 항이 $f(p_i)$ 하나임을 이용했다. $F_{prime}$ 표현을 이용하면, $$ F_k(n) = \left( \sum_{k \le i, p_i^2 \le n} \sum_{c \ge 1, p_i^{c+1} \le n} \left( f(p_i^{c+1}) + f(p_i^c) F_{i+1}(n/p_i^c) \right) \right) + F_{prime}(n) - F_{prime}(p_{k-1})$$을 얻는다. 이제 이 점화식을 써먹는 방법을 고민해보자.

  • 위 점화식을 top-down으로 계산한다고 가정하자.
  • 또 $\lfloor n/i \rfloor$ 형태의 값들만 등장한다. 이제 $\mathcal{O}(\sqrt{n})$개의 값들만 등장함은 익숙하다.
  • 재귀 $k$ 값들을 보면, $p_i^2 \le n$을 아는 상태에서 $F_{i+1}(*)$ 값들을 호출함을 볼 수 있다. 
  • 그러니 $F_k(n)$을 구해야 한다는 재귀호출이 나왔다면, $p_{k-1}$이 큰 값은 아니라는 것이다. 
  • 왜 이게 중요하냐면, $F_{prime}(p_{k-1})$의 값이 전처리의 영역에 있다는 것을 설명하기 때문이다. 
  • 즉, $n^{1/2}$ 정도의 값까지 $F_{prime}$의 값을 전처리하면, $F_{prime}(p_{k-1})$ 값은 $\mathcal{O}(1)$에 정리된다.
  • 이제 $\lfloor n/i \rfloor$ 형태의 값들에 대해 $F_{prime}$ 값을 전부 구해야 한다는 문제가 있다.
  • 만약 $f(p)$가 $p$에 대한 다항식이라면, 앞에서 배운 Lucy-Hedgehog 알고리즘으로 이를 전부 전처리할 수 있다.

시간복잡도 분석은 자료에 따르면 $\mathcal{O}(n^{1-\epsilon})$이다. 하지만 시간복잡도에 비해 매우 좋은 성능을 보인다.

이 알고리즘의 강점은 xudyh's sieve처럼 $g$를 고민해서 잡을 필요가 없다는 점, 그리고 $f$에 대해 가정하는 점이 적다는 점이다.


구현 노트.

  • 중요한 사실인데, $p_i < \sqrt{n} < p_{i+1}$이라고 해서 $n/p_i < p_{i+1}$인 것은 아니다.
  • 그러니 전처리를 하는 소수/체의 범위를 조금 조심해서 고를 필요가 있다. 마음 편하게 $\sqrt{n} + 500$ 정도로 두자.


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


이 글에서는 다음 알고리즘을 다룬다.

  • Lucy-Hedgehog의 Algorithm 소개 및 최적화.
  • Meissel-Lehmer Algorithm 소개 및 최적화.


Lucy-Hedgehog Algorithm

  • 글을 읽기 전에, 경건한 마음을 가지고 Project Euler에 가입하자. 그 다음, (쉬운) 10번 문제를 풀자.
  • 그러면 여러 사람이 자신의 접근과 풀이, 코드를 공유하는 thread에 접근할 수 있다.
  • 이제 https://projecteuler.net/thread=10;page=5#111677로 가면, 이 알고리즘의 기원을 볼 수 있다.
  • Note. 동적계획법만 알고 있어도 이 알고리즘을 충분히 이해할 수 있다.


$n$ 이하의 소수의 개수를 $\mathcal{O}(n^{3/4})$에 계산하는 방법을 소개한다. 먼저, $n^{1/2}$ 이하의 소수들을 에라토스테네스의 체를 이용하여 전처리한다. 

이제 $dp[v][m]$을 에라토스테네스 체에서 $m$까지 보았을 때, $1$ 이상 $v$ 이하 자연수 중 지워지지 않은 것의 개수라고 하자.

  • 즉, $dp[v][m]$은 소수거나, 모든 소인수가 $m$ 초과인 $1$ 이상 $v$ 이하 자연수의 개수이다.
  • 특히, $(m+1)^2 > v$라면 모든 소인수가 $m$ 초과인 $v$ 이하 자연수는 무조건 소수이므로, $dp[v][m]$은 $v$ 이하 소수의 개수다.


이제 $dp$ 식을 생각해보자.

  • 만약 $m$이 소수가 아니라면, 애초에 에라토스테네스 체에서 하는 일이 없다. 이때는 $dp[v][m]=dp[v][m-1]$.
  • 만약 $m^2 > v$라면, $m$이 소수여도 이번에 새로 지워지는 값은 없을 것이다. 여전히 $dp[v][m]=dp[v][m-1]$.
  • $m$이 소수라고 가정하자. 이번에 제거되는 값들은 최소 소인수가 $m$인 값들이다.
  • 즉, $m$과 "최소 소인수가 $m-1$ 초과인 자연수"을 곱한 값들이 지워진다. 
  • 이 값의 개수를 구하려면, $\lfloor v/m \rfloor$ 이하 자연수 중 최소 소인수가 $m-1$ 초과인 자연수의 개수를 구하면 된다.
  • 그런데 $dp[\lfloor v/m \rfloor][m-1]$에 소수이거나, 최소 소인수가 $m-1$ 초과인 값의 개수가 저장되어 있다.
  • 이 값에 $m-1$ 이하 소수의 개수를 빼면 원하는 값이 나올 것이다. 그 값은 $dp[m-1][m-1]$에 저장되어 있다.
  • 그러므로, $dp[v][m] = dp[v][m-1] - (dp[\lfloor v/m \rfloor][m-1] - dp[m-1][m-1])$이라는 결과가 나온다.


이제 필요한 $dp$ 상태를 생각해보자. 우리는 $dp[n][\lfloor \sqrt{n} \rfloor]$가 필요하다.

  • 먼저 두 번째 인자 $m$을 보자. 우리의 관찰에 의해, $m^2 > v$인 경우는 추가적으로 고려할 필요가 없다. 
  • 특히, 우리가 다룰 모든 두 번째 인자 $m$ 값은 $\sqrt{n}$ 이하이다.
  • 이제 첫 인자인 $n$을 생각하자. 필요한 첫 번째 인자의 값은 $\lfloor n/i \rfloor$ 형태의 값을 가짐을 알 수 있다.
  • 즉, 필요한 첫 번째 인자의 값은 $1, 2, \cdots, \lfloor \sqrt{n} \rfloor$과 $\lfloor n/1 \rfloor, \lfloor n/2 \rfloor, \cdots , \lfloor n/\lfloor \sqrt{n} \rfloor \rfloor$다.
  • 중요한 점은, 다루어야 할 첫 번째 인자의 값의 종류가 $\mathcal{O}(\sqrt{n})$개라는 점이다. 이제 준비 완료.


이제 본격적으로 알고리즘을 구체화한다.

  • DP에 익숙한 사람들은 바로 눈치를 챘겠지만, 위 DP 식은 정말 토글링을 하기 좋은 식이다. 
  • 그러니 $dp$를 1차원 배열로 정의하고, 순서대로 $m$ 값을 늘려나가보자. $dp$ 배열은 $dp[v][m]$을 갖고 있을 것이다.
  • 그러기 위해서는 사실 $dp[v][1]$을 정의해야 하는데, 이는 시작하자마자 체에서 1을 제거했다고 치고 $dp[v][1]=v-1$로 둔다.
  • 이제 큰 $v$값부터 순서대로 $dp[v][m]=dp[v][m-1]-(dp[\lfloor v/m \rfloor][m-1] - dp[m-1][m-1])$을 적용한다.
  • 그러다가, $v \le m^2$인 시점이 오면 업데이트를 중단하고 다음 $m$으로 넘어가도록 한다. 
  • 물론, 우리가 실제로 업데이트를 해야 하는 $v$의 값들은 앞서 언급한 $\mathcal{O}(\sqrt{n})$ 개의 값들이다.
  • 마지막으로 얻은 $dp[n]$의 값이 우리가 원하는 $n$ 이하의 소수 개수다. 이게 알고리즘 끝이다. 예상보다 간단했을 것.
  • 특히, 실제로 관심이 있던 각 값 $v$에 대하여, $dp[v]$는 $v$ 이하의 소수 개수다. 즉, 예를 들어, $n/2$, $n/3$ 이하의 소수 개수도 동시에 구해진다.


시간복잡도 분석을 하자.

  • 이 알고리즘은 "실제로 계산해야 하는" $dp[v][m]$의 개수만큼 계산을 한다.
  • 특히, 우리는 $m^2 \le v$인 경우에만 추가적인 계산을 하니, 우리가 다루는 각 $v$값에 대해 $\sqrt{v}$를 더한만큼 계산을 한다. 이는 $$ \sum_{i=1}^{\sqrt{n}} \sqrt{i} + \sum_{i=1}^{\sqrt{n}} \sqrt{n/i} = \mathcal{O}(\sqrt{n} \cdot \sqrt{n}^{1/2}) + \mathcal{O}(\sqrt{n} \cdot \sqrt{n}^{1/2}) = \mathcal{O}(n^{3/4})$$다. 여기서 $1+1/\sqrt{2} + 1/\sqrt{3} + \cdots + 1/\sqrt{n} = \mathcal{O}(\sqrt{n})$을 이용하였다.


구현 노트.

  • $dp$ 배열을 관리하기 위하여 가장 쉽게 생각할 수 있는 방법은 역시 C++의 std::map일 것이다. 하지만 로그가 붙는다.
  • 로그를 붙이지 않고 구현을 하려면, $dp$를 말 그대로 배열에 저장해야 한다. 이를 위해서는, $v = 1, 2, \cdots, \lfloor \sqrt{n} \rfloor$에 대응되는 길이 $\lfloor \sqrt{n} \rfloor$의 배열과, $v = \lfloor n/1 \rfloor, \lfloor n/2 \rfloor, \cdots ,  \lfloor n/\lfloor \sqrt{n} \rfloor \rfloor$에 대응되는 길이 $\lfloor \sqrt{n} \rfloor$의 배열을 따로 두면 된다. 자세한 디테일은 github 참고.


추가 코멘트.

  • 필자는 이 알고리즘을 많이 써먹었다. 이 알고리즘의 장점은, 확장이 쉽게 가능하다는 것이다.
  • 예를 들어, 소수의 개수가 아니라 소수의 합, 소수의 제곱의 합 등을 원하는 경우도 $dp$ 식만 살짝 바꿔주면 해결된다.
  • 심지어, $n$ 이하 $4k+1$ 꼴 소수의 개수/합, $4k+3$ 꼴 소수의 개수/합 역시 약간의 아이디어를 추가하면 할 수 있다.
  • 확장이 쉽게 가능한 이유는, 역시 기본 원리인 dp 자체가 간단한 편이기 때문이다. 이해하기도 쉬운 편인 알고리즘이라 생각. 
  • 사실 시간복잡도 분석을 빡빡하게 하지 않았다. 실제로는 $m$이 소수인 경우만 중요하다는 사실도 역시 써먹을 수 있다.
  • 빡빡한 시간복잡도 분석을 위해서는, $n$ 이하 소수의 개수가 약 $n/\ln n$개 존재한다는 Prime Number Theorem을 참고하라.


이 알고리즘을 $\mathcal{O}(n^{2/3})$ 수준으로 최적화시킬 수 있다. 여기서 "로그가 시간복잡도에 얼마나 붙는지"는 따지지 않겠다.

  • 지수를 3/4에서 2/3으로 내리는 것은 11단원에서도 등장한다. 보통은, "에라토스테네스의 체를 더 써먹자"의 방식으로 최적화가 된다.
  • 사실 우리는 기존 알고리즘에서 에라토스테네스의 체로 $1$부터 $n^{1/2}$까지의 소수 판별만 했다. 체를 더 써먹자.

이제부터, 기존 점화식 $dp[v][m] = dp[v][m-1] - (dp[\lfloor v/m \rfloor][m-1] - dp[m-1][m-1])$을 top-down으로 돌려보자.

목표는 실제로 계산을 전부 하는 것이 아니라, 어떤 $dp$ 값들이 필요한지를 탐색하는 것이다. 시작은 물론 $dp[n][\lfloor \sqrt{n} \rfloor]$을 호출하는 것이다.

  • $dp[m-1][m-1]$은 $m$ 미만의 소수 개수와 같다. 이는 에라토스테네스의 체로 처리할 수 있으니, $dp$ 값으로 간주하지 않는다.
  • 만약 $m=1$이라면, $dp[v][m]=v-1$임을 이미 알고 있으니, 이를 필요한 $dp$ 값으로 간주하지 않는다.
  • 만약 $v \le n^{2/3}$이라면, "$dp[v][m]$을 계산해야 한다"라는 정보를 저장하고 리턴한다.
  • 재귀 과정에서 이미 방문한 위치 $v, m$에 다시 도달했다면, 이미 $dp[v][m]$을 계산하기 위해 필요한 $dp$ 위치를 알고 있는 상태이므로 바로 리턴.
  • 즉, $(v,m)$의 자손이 $(v,m-1)$과 $(\lfloor v/m \rfloor ,m-1)$인 이진트리를 생각하고, $m=1$이거나 $v \le n^{2/3}$인 경우를 리프 노드로 둔다.


시간복잡도를 분석하자.

  • 사실 1. 위 재귀 루틴의 시간복잡도는 $\mathcal{O}(n^{2/3})$이다. 
  • 증명. $dp$ 상태 중, $v \ge n^{2/3}$인 것의 개수를 구하면 충분하다. 실제로 관심 있는 $v$만 고려하면, 이는 $$ \sum_{i=1}^{n^{1/3}} \sqrt{n/i} = \mathcal{O}(\sqrt{n} \cdot \sqrt{n^{1/3}}) = \mathcal{O}(n^{2/3})$$
  • 사실 2. 결과적으로 우리가 "계산해야 하는 $dp$ 값"들의 위치는 $\mathcal{O}(n^{2/3})$개이다.
  • 증명. 물론, 이는 사실 1에서 바로 따라온다. 애초에 시간복잡도가 저러니 저장한 위치의 개수도 마찬가지.

이제 우리는 $\mathcal{O}(n^{2/3})$개의 $dp$ 값만 구해주면 된다. 이 값들을 다 구하면, 다시 재귀를 돌려 답을 얻을 수 있다.

즉, 이전에 "$dp[v][m]$을 계산해야 한다"고 하고 리턴한 경우, 이번에는 실제로 계산된 $dp[v][m]$의 값을 리턴하면 된다!

이제부터 우리는 알고리즘의 영역으로 온다. 처음으로 이 가이드에서 C++ STL에 없는 자료구조를 사용한다.

  • 오프라인으로 $dp$ 값들을 계산한다. $m$의 값이 증가하는 순서로 필요한 값들을 계산할 것이다.
  • $dp[v][m]$을 계산하는 것은, 결국 정의상 $m$까지 체를 돌렸을 때 $1$부터 $v$까지 살아남은 자연수의 개수를 구하는 것이다.
  • 현재 자연수가 남아있는지 여부를 0/1로 표현하면, 이는 결국 $[1, v]$에서 구간합을 구하는 것과 마찬가지다.
  • 에라토스테네스의 체에서 값을 지우는 것은 결국 해당 자연수에 대응되는 값을 1에서 0으로 바꾸는 것과 같다.
  • 그러니 이 문제는 사실 Point Update + Segment Query. 세그먼트 트리로 해결할 수 있다. (속도를 위해 펜윅을 쓰는 것을 추천한다.)

이제 전체 문제가 $\mathcal{O}(n^{2/3})$ 시간에 해결된다. 필요한 에라토스테네스의 체 크기가 $n^{2/3}$이었음에 주목하자.

  • 이 최적화 방법은 앞서 언급한 Lucy-Hedgehog 알고리즘의 여러 확장에도 바로 적용된다.
  • 예를 들어, 소수의 합을 구하고 싶다면 자연수 $k$가 남아있는지 여부를 0/$k$로 나타내면 된다.


Meissel-Lehmer Algorithm

  • 사실 Lucy-Hedgehog와 큰 차이는 없다. 개인적으로 많이 사용하지 않아, 얼마나 확장이 쉽게 되는지는 잘 모르겠다.
  • 하지만 알고리즘의 원리가 근본적으로 비슷하니, 아마 크게 다르지는 않을 것이라고 생각한다. 독자에게 맡긴다.


이번에는 notation을 살짝 바꾼다. 그래도 Lucy-Hedgehog 알고리즘과 비슷한 점이 많으니, 비교해가며 읽자.

  • $p_k$를 $k$번째 소수로 정의한다. 단, $1$-index를 이용한다. 즉, $p_1, p_2, p_3 = 2, 3, 5$.
  • $\pi(x)$를 $x$ 이하의 소수의 개수로 정의한다. 이는 매우 standard 한 notation이다. 예를 들어, $\pi(x) \approx x / \ln x$.
  • $\phi(x, a)$를 모든 소인수가 $p_{a+1}$ 이상인 $x$ 이하 자연수 개수라 하자. 중요한 점은, $p_1, p_2, \cdots ,p_a$ 역시 제거했다는 것이다.
  • 또한, $\phi(x, a)$를 생각할 때 $1$도 포함한다. 그러니, 정확하게 말하자면 $p_1, p_2, \cdots , p_a$의 배수를 제거한 결과를 보는 것이다.
  • 이어서, $P_k(x, a)$를 모든 소인수가 $p_{a+1}$ 이상인 $x$ 이하 자연수 중, 소인수 개수가 정확히 $k$개인 것의 개수라 하자.
  • 이때, "소인수 개수"란, $n=p_1^{e_1} \cdots p_k^{e_k}$라 할 때 $e_1 + e_2 + \cdots + e_k$로 정의된다. 
  • 특히, $p_{a+1}^k > x$라면 $P_k(x, a) = 0$이 성립한다. 이제 필요한 용어의 정의는 끝났다.


$\pi(x)$를 계산하려면, 필요한 값은

  • 우선 억울하게 제거당한 소수들 $p_1, p_2, \cdots , p_a$가 있다. 총 $a$개.
  • 이제 모든 소인수가 $p_{a+1}$ 이상인 $x$ 이하 자연수가 필요하다. 총 $\phi(x, a)$개.
  • 그런데 $1$은 소수가 아니니, $1$을 빼주어야 할 것이다.
  • 또한, 소인수 개수가 $2$ 이상인 것에는 관심이 없으니, $P_2(x, a) + P_3(x,a) + \cdots $를 빼야한다. 즉, $$ \pi(x) = a + \phi(x, a) - 1 - \sum_{k=2}^\infty P_k(x, a)$$
  • 물론, $p_{a+1}^k > x$면 $P_k(x, a)=0$이므로, 실제로는 무한합을 계산할 필요가 없다. 


특히, 우리는 $a = \pi(x^{1/3})$을 선택할 것이다. 그러면 $p_{a+1}^3 > x$이므로, $P_3$부터 쓸모가 없다.

그러므로, $\pi(x) = a + \phi(x, a) - 1 - P_2(x, a)$를 계산하면 된다. 이제 각 항을 계산할 전략을 설계해보자.

  • $a$의 계산. 단순히 $x^{1/3}$ 이하의 소수의 개수이므로, 에라토스테네스의 체로 정리된다.
  • $\phi(x, a)$의 계산. 이 부분은 Lucy-Hedgehog 알고리즘의 dp와 크게 다르지 않다.
  • $\phi(x, a)$를 계산해보자. 일단 $\phi(x, a-1)$에서 시작하면, $a-1$에서 $a$로 넘어갈 때 새로 제거되는 수의 개수를 세면 된다.
  • 이 값을 계산하려면, "최소 소인수가 정확히 $p_a$인 $x$ 이하의 자연수"의 개수를 세면 된다.
  • 이는 결국 "모든 소인수가 $p_{a}$ 이상인 $\lfloor x/p_a \rfloor$ 이하의 자연수"의 개수와 같다. 이는 $\phi(\lfloor x/p_a \rfloor, a-1)$.
  • 그러므로, dp 식 $\phi(x, a) = \phi(x, a-1) - \phi(\lfloor x/p_a \rfloor, a-1)$이 성립함을 알 수 있다.
  • 우선 Lucy-Hedgehog에서 본 DP와 여러 성질이 겹친다. $\lfloor x/i \rfloor$ 형태의 값만 필요하다는 점 역시 똑같다.
  • 토글링이 가능한 DP 형태라는 점도 똑같다. 그러니 이미 살펴본 접근 방식과 거의 동일한 방법이 여기서도 먹힌다.
  • 하지만 이 방법의 중요한 점은, $p_a^2 > x$여도 $\phi(x, a)$와 $\phi(x, a-1)$이 다르다는 것이다. 
  • 에라토스테네스의 체와 다르게 여기서는 $p_a$도 제거하므로, $\phi(x, a) = \phi(x, a-1) - 1$이 성립하겠다.
  • 물론, $x < p_a$면 $\phi(x, a) = \phi(x, a-1)$이 성립함은 당연하다.
  • 그래서 기존 $\mathcal{O}(n^{3/4})$ 알고리즘을 적용하기가 조금 귀찮아진다. $1$을 진짜 정직하게 빼주면 당연히 시간복잡도가 유지되지 않는다.
  • 그러니 $p_a^2 > x$면 $\phi(x, a)$를 업데이트 하지 않고, 대신 "실제로 $1$이 몇 번 더 빠졌어야 하는가"를 필요할 때 계산해주자.
  • 재밌는 점은 $\mathcal{O}(n^{2/3})$ 최적화 알고리즘은 그대로 적용된다는 것이다.
  • 언급한 사실 1, 사실 2 역시 그대로 유지되며, 오프라인 쿼리 + 세그트리/펜윅트리로 $\phi$를 계산할 수 있음도 동일하다.
  • $P_2(x, a)$의 계산. 이 값은 다행히도, 정직하게 계산할 수 있다.
  • $P_2(x, a)$는 결국 $a<b \le c$를 만족하면서 $p_bp_c \le x$인 $(b, c)$의 순서쌍의 개수를 구하는 것이다.
  • $b$를 고정하고, $c$의 개수를 세자. 그러면 $p_c \le x/p_b$와 $b \le c$가 동시에 성립해야 하니, $\pi(\lfloor x/p_b \rfloor) - (b-1)$이다.
  • 가능한 $b$의 범위는 물론 $a+1$부터 $\pi(x^{1/2})$까지다. 즉, 이 결과를 합치면 우리는 $$P_2(x, a) = \sum_{b=a+1}^{\pi(x^{1/2})} \left( \pi(\lfloor x/p_b \rfloor) - (b-1) \right)$$를 얻는다. 이제 이 값을 빠르게 계산하면 된다. 여기서 핵심은 주어진 범위의 $b$에 대하여, $x/p_b < x^{2/3}$이 성립한다는 것이다. 
  • 그러니 실제로 필요한 것은 $x^{2/3}$까지의 에라토스테네스의 체 결과다. 그러니, 모든 계산을 $\mathcal{O}(n^{2/3})$ 시간에 할 수 있다.


이를 종합하면 우리가 원하는 알고리즘이 나온다. 최적화로 $\mathcal{O}(n^{2/3})$까지 가는 방법도 설명하였다. 그런데,

  • 소수 카운팅 알고리즘의 속도를 측정할 수 있는 가장 쉬운 방법은, Codeforces의 Four Divisors 문제를 해결하는 것이다.
  • 당장 Editorial에서도 이 글에서 다룬 $\mathcal{O}(n^{2/3})$ 테크닉을 언급한다. 자세한 설명이 있으니 읽어보는 것도 좋다.
  • 그런데, 정작 submission을 실행 속도로 정렬한 다음 몇 개의 제출을 읽어보면 오프라인이고 세그트리고 다 없다.
  • 즉, 시간복잡도로 표현하는 속도와 실제 실행 속도와의 괴리감이 약간 있는 부분이라고 볼 수 있겠다. 이제 실제로 빠른 알고리즘을 간략하게 설명한다.
  • (물론, 효율적인 코드를 매우 잘 작성하는 분들이라면 $\mathcal{O}(n^{2/3})$ 알고리즘으로 저 분들을 이길 수도 있을 것이다.)
  • 먼저, 에라토스테네스의 체를 많이 돌린다. $n=10^{11}$ 기준으로, 500만 정도 크기로 돌린 것 같다.
  • 앞서 우리는 $\pi(x^{1/3})$에서 식을 자르는 방식으로 $P_3(x, a)$ 항을 날렸다. 이번에는 $\pi(x^{1/4})$에서 식을 잘라서 $P_3(x, a)$ 항을 살린다.
  • 즉, $a = \pi(x^{1/4})$를 잡고, $\pi(x) = a + \phi(x, a) - 1 - P_2(x, a) - P_3(x, a)$를 계산한다.
  • 이제 추가적으로 $P_3$에 대한 전개식을 계산해야 한다. $p_bp_cp_d$ 형태를 보면 되는데, 진짜 정직하게 $b, c$에 대한 이중 for 문을 돌린다.
  • 이제 $\phi(x, a)$를 계산해야 한다. 이 값을 계산하기 위해서, 상당한 최적화를 한다.
  • $\phi(x, a)$의 context에서, 특정 값이 살아있는지 여부는 그 값을 $p_1p_2 \cdots p_a$로 나눈 나머지에 의해 결정된다.
  • 그러니, $a$가 작은 경우 ($a \le 7$을 이용한다) 아예 어떤 나머지가 살아남는지를 통째로 전처리를 한다. 크기 $p_1 \cdots p_a$짜리 전처리다.
  • 이러면 $a \le 7$인 입력이 들어온 경우, 아예 $\mathcal{O}(1)$에 $\phi(x, a)$를 반환할 수 있다. 이것이 첫 최적화.
  • 만약 $x < p_a^2$이라면, $\phi(x, a)$는 사실 $p_{a+1}$ 이상 $x$ 이하 소수의 개수이다. 
  • 그러니, 만약 $x$ 이하의 소수 개수가 체로 전처리된 상황이라면, 바로 답을 반환할 수 있다. 이것이 두 번째 최적화.
  • 만약 $x < p_a^3$이라면, $\phi(x, a) = \pi(x) - a + 1 + P_2(x, a)$다. 이때, $\pi(x)$의 값이 이미 체로 전처리된 상황이라 하자.
  • 이 경우에는, $P_2(x, a)$를 기존 방식과 동일한 방법으로 구해 $\phi(x, a)$를 얻는다. 이것이 세 번째 최적화.
  • 만약 이 모든 경우에 해당하지 않는다면, 그제서야 점화식 $\phi(x, a) = \phi(x, a-1) - \phi(\lfloor x/p_a \rfloor, a-1)$을 호출한다.
  • 사실 기존 Meissel-Lehmer 알고리즘에서 $a = \pi (x^{1/2})$를 선택하는 것도 방법이다. 방법은 정말 많다..
실제로 어떻게 구현되었는지 보려면, 이 링크를 타고 여러 코드를 탐색해보자. 상당히 빠르다 :) 

$\mathcal{O}(n^{2/3})$ 알고리즘들이 메모리를 많이 잡아먹는데 (물론, 적게 먹는 variant가 있는 것으로 안다) 주의해야 한다.

필자는 가능하면 Lucy-Hedgehog를 쓰고 (특히 Project Euler와 같은 시간제한이 없는 경우에는 더더욱) 만약 시간제한이 빡빡하다면 Meissel-Lehmer의 최적화된 구현체, 즉 위 Four Divisors 문제의 링크에 있는 알고리즘을 사용하여 문제를 푼다. 모두 흥미로운 알고리즘들이다. 이 정도로 이번 내용은 끝.

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는 정확히 같은 알고리즘이다. 실제 코드도 별로 다를 것이 없다.