Discrete Log（離散対数問題） 〜Quantum Zooやっていく〜

July 19, 2021, 3:14 p.m. edited Feb. 11, 2023, 4:04 p.m.

#Quantum Zoo  #量子情報  #Qiskit  #量子力学 

$$ \def\bra#1{\mathinner{\left\langle{#1}\right|}} \def\ket#1{\mathinner{\left|{#1}\right\rangle}} \def\braket#1#2{\mathinner{\left\langle{#1}\middle|#2\right\rangle}} $$

まさかの第二回。今回は離散対数問題を解く Shor のアルゴリズムを扱う。

非推奨の項目がやたら多いが、そこはまだあまりわかってない頃に書いたものなので無視してほしい。

背景

P. Shor が素因数分解アルゴリズムを公表した論文内で素因数分解アルゴリズムとともに示されていたもの。素因数分解と同様に暗号と強い関わりが存在する。具体的な問題内容は次節の「問題」で示す。

最速の古典アルゴリズムでは \(\exp(O(\log p)^{1/3}(\log\log p)^{2/3})\) （\(p\) の定義は次節の「問題」にて）の時間計算量がかかる一方、 Shor による量子アルゴリズムでは \(\mathrm{poly}(\log p)\) の時間計算量で解ける。

問題

\(a,b,p\in\mathbb{Z}\) (\(p\) は素数でないと古典でもあまり難しくないので、今回は素数の場合を考える) が与えられたときに

$$a^s\equiv b\pmod p$$

を満たす最小の正の整数 \(s\) を求める。

アルゴリズム

素因数分解のときと同じく位数発見アルゴリズムが重要な役割を果たすが、若干より複雑になっている。このアルゴリズムは例によって量子コンピュータと量子通信 II に途中（連分数展開によりどのように最終的な結果を得るかは載ってない）まで載っており、位数発見アルゴリズムのときと同様に詳しくはそちらを読んでほしい。

1. 位数発見アルゴリズムを用いて、 \(p\) を法とした \(a\) の位数 \(r\) （つまり \(a^r\equiv 1\pmod p\)）を見つける
2. \(t=\lceil\log p\rceil\) として、 \(t\) 個の量子ビットからなる第一レジスタ、第二レジスタ、第三レジスタを用意： \(\ket{0^t}\ket{0^t}\ket{0^t}\)
3. アダマール演算により重ね合わせを作る： \(\frac{1}{2^t}\sum_{x_1=0}^{2^t-1}\sum_{x_2=0}^{2^t-1}\ket{x_1}\ket{x_2}\ket{0^t}\)
4. \(f(x_1,x_2)=b^{x_1}a^{x_2}\) として、冪剰余計算により次の状態を得る：
   $$\frac{1}{2^t}\sum_{x_1=0}^{2^t-1}\sum_{x_2=0}^{2^t-1}\ket{x_1}\ket{x_2}\ket{f(x_1,x_2)}$$
5. \(\hat{f}(l_1,l_2)=\frac{1}{\sqrt{r}}\sum_{j=0}^{r-1}e^{-2\pi il_2j/r}\ket{f(0,j)}\) とすると、 4. の状態は次のように式変形される：
   $$\frac{1}{2^t\sqrt{r}}\sum_{l=0}^{r-1}\sum_{x_1=0}^{2^t-1}\sum_{x_2=0}^{2^t-1}e^{2\pi i(slx_1+lx_2)/r}\ket{x_1}\ket{x_2}\ket{\hat{f}(sl,l)}$$
6. さらに次のように式変形される：
   $$\frac{1}{2^t\sqrt{r}}\sum_{l=0}^{r-1}\left[\sum_{x_1=0}^{2^t-1}e^{2\pi i(slx_1)/r}\ket{x_1}\right]\left[\sum_{x_2=0}^{2^t-1}e^{2\pi i(lx_2)/r}\ket{x_2}\right]$$
7. 第一、第二レジスタに量子フーリエ逆変換：
   $$\frac{1}{\sqrt{r}}\sum_{l=0}^{r-1}\ket{\widetilde{sl/r}}\ket{\widetilde{l/r}}\ket{\hat{f}(sl,l)}$$
8. 第一、第二レジスタを測定すると、 \(\widetilde{sl/r}\), \(\widetilde{l/r}\) を得る

連分数展開などにより \(s\) を得る

位数発見アルゴリズムのときと同じく、やはり \(s\) を得るには古典アルゴリズムによる後処理が必要となる。今回はさらに第一、第二レジスタ 2 つからの結果があるので事態をややこしくしている。そのうえ量子コンピュータと量子通信 II にはその方法は載っていないので自力で考えたものをここに載せる。

連分数展開により \(l\) を得る

これは簡単。位数発見アルゴリズムのときの連分数展開により \(r\) を求めたときと同じ。それどころか、今回は既に \(r\) がわかっているので、分母に \(r\) が出てきたら分子が正しい \(l\) となる。なお、 \(r\) が素数でない場合は \(r\) の因数が分母に出てくることがあるが、のちの都合により \(l\) と \(r\) が互いに素であってほしいので、 \(r\) そのものが出てこなかったら量子アルゴリズムのやり直しとすることにする。

連分数展開により \(\beta\) を得る

\(l\) と同じように \(sl\) も求まりそうなものだが、 \(sl/r>1\) となりうるためそうはいかない。つまり、 \(sl/r=\alpha+\beta/r\) (\(\alpha,\beta\in\mathbb{Z}\)) とすると、測定で得られるのは \(\widetilde{\beta/r}\) のみになるということである。ゆえに、まずは先ほどと同じ方法で正しい \(\beta\) を得る。

拡張ユークリッド互除法により \(s\) を得る

$$ \begin{align} \frac{sl}{r}&=\alpha+\frac{\beta}{r}\\ ls-r\alpha&=\beta\tag{1} \end{align} $$

最後に式 (1) の中で未知のものとして残っているのは \(s\) と \(\alpha\) のみである。ここで、 \(l\) と \(r\) は互いに素であるので、拡張ユークリッド互除法を適用して \(s\) と \(\alpha\) の整数解を得ることができる。あとは \(0\leq s<r\) の条件を満たすものが答えとなる。

Qiskit による実装

位数発見アルゴリズムのときと同様に、冪剰余計算には elementary.py を用いている。以下に Qiskit による実装を載せる。コードの全体は GitHub に置いてある。

from typing import Optional
import math

import numpy as np
import matplotlib.pyplot as plt
from sympy import gcdex
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, Aer, transpile, assemble
from qiskit.visualization import plot_histogram
from sympy import Rational, gcdex
from sympy.ntheory.continued_fraction import continued_fraction, continued_fraction_convergents

from qft import qft
from elementary import ax_modM
from order_finding import order_finding
from classical_utils import decode_bin

def discrete_log(a: int, b: int, p: int, show_hist: Optional[bool] = False) -> int:
    """Shor's discrete log algorithm: given $a,b,p\in\mathbb{Z}$, it finds $s$ such that $a^s\equiv b\pmod p$.
    Args:
        a (int): $a$
        b (int): $b$
        p (int): $p$
    """

    r = order_finding(x=a, N=p, show_hist=False)
    t = int(np.ceil(np.log2(p)))

    first_register = QuantumRegister(t)
    second_register = QuantumRegister(t)
    third_register = QuantumRegister(2 * t)
    auxiliary_register_mid = QuantumRegister(t)
    auxiliary_register_end = QuantumRegister(6 * t - 2)
    classical_register = ClassicalRegister(2 * t)

    qc = QuantumCircuit(
        first_register,
        second_register,
        third_register,
        auxiliary_register_mid,
        auxiliary_register_end,
        classical_register,
    )

    qc.h(first_register)
    qc.h(second_register)

    qc.append(ax_modM(a=b, M=p, N_len=t), list(first_register) + list(auxiliary_register_mid) + list(third_register) + list(auxiliary_register_end))
    qc.append(ax_modM(a=a, M=p, N_len=t, x_0_at_first=False), list(second_register) + list(auxiliary_register_mid) + list(third_register) + list(auxiliary_register_end))

    qc.append(qft(n=t).inverse(), first_register)
    qc.append(qft(n=t).inverse(), second_register)

    qc.measure(list(first_register) + list(second_register), classical_register)
    #qc.measure(third_register, classical_register)

    backend = Aer.get_backend('aer_simulator_matrix_product_state')#('aer_simulator')
    qc = transpile(qc, backend)
    job = backend.run(qc, shots=10000)
    hist = job.result().get_counts()

    if show_hist:
        figsize_x = max(7 * (len(hist) // 8), 7)
        plot_histogram(hist, figsize=(figsize_x, 5))
        plt.savefig(f'img/discrete_log_a{a}_b{b}_p{p}_r{r}.png', bbox_inches='tight')

    for measured_key, _ in sorted(hist.items(), key=lambda x: x[1], reverse=True):
        tilde_l_per_r = Rational(decode_bin(measured_key[:t]), 2 ** t) # decoded from second register: $\widetilde{l/r}$
        if tilde_l_per_r == 0:
            continue

        l = None
        for fraction in continued_fraction_convergents(continued_fraction(tilde_l_per_r)):
            if fraction.denominator == r:
                l = fraction.numerator # get correct $l$
                break

        if l is None:
            continue

        tilde_beta_per_r = Rational(decode_bin(measured_key[-t:]), 2 ** t) # decoded from first register: $\widetilde{\beta/r}$
        if tilde_beta_per_r == 0:
            continue

        beta = None
        for fraction in continued_fraction_convergents(continued_fraction(tilde_beta_per_r)):
            if fraction.denominator == r:
                beta = fraction.numerator # get correct $\beta$
                break

        if beta is None:
            continue

        s, alpha, _ = gcdex(l, -r)
        s = int(s * beta)
        if pow(a, s, p) == b:
            return s

    raise Exception('s is NOT found!')


if __name__ == '__main__':
    print(discrete_log(a=2, b=4, p=7, show_hist=True))

例として、 \(a=2\), \(b=4\), \(p=7\) で実行する（\(t=3\)）と、まず位数 \(r=3\) が得られたうえで以下のようなプロットが得られる：

0 のビンは無視するとして、確率の大きいビンである \(011101=\ket{101}\ket{011}=\ket{5}\ket{3}\) に着目する。 \(2^t=8\) より、ここから得られる \(\widetilde{l/r}\) は

$$\widetilde{l/r}=\frac{5}{8}=\frac{1}{1+\frac{1}{1+\frac{1}{1+\frac{1}{2}}}}$$

より、連分数展開により

$$\frac{1}{1+\frac{1}{1+\frac{1}{1}}}=\frac{2}{3}=\frac{l}{r}$$

から \(l=2\) となる。同様に、 \(\frac{3}{8}\) からの連分数展開により \(\beta=1\) と得られる。よって、

$$2s-3\alpha=1$$

より、 \(s=-1\), \(\alpha=-1\) と求まる。この結果は、もう一本の大きい確率をもつビンである \(101011\) からも同じく求まる。

あとは、モジュラ逆数の方が手に入っていたことから \(r+s=2\) が求めている答えとなる。実際に \(a^2\equiv b\pmod p\) が成り立つ。以上。

IPA 文書について補足

Google で「離散対数 量子」でググると結構上に出てくる IPA の「量子計算機の研究動向に関する調査」 調査報告書 のうちの 4.3 暗号解読に関係するアルゴリズム で紹介されている離散対数の量子アルゴリズムは結構シンプルで良さそうである。ただし、注意書きにあるように、この IPA 文書で紹介されているアルゴリズムでは \(p-1\) が smooth である必要がある。

ここでの smooth とは、 \(p-1\) が 2 の累乗で表されるということをおそらく意味している。これにより、量子フーリエ変換を \(p-1\) を分母にして適用できるため、ここまでシンプルなものになるようである。（参考：量子フーリエ変換が理由である説明）

ゆえに、一般の、古典では難しい場合も含めたものについては、前節までに記した量子アルゴリズムを使えば良いということになる。

以下、非推奨

2023-02-02 追記：以下はまだあまりわかってない頃に書いた内容なので、↑までの正しい（はず）の内容だけを見ていれば良い

なお、いきなり結論に向かいたい読者は アルゴリズム（測定を最後にした場合） および 測定で得られた整数値から \(d\) を計算（おそらく完全版）、実装も見る場合は Qiskit による実装 を見れば良い。

では、なぜこのような構成になっているかというと、本記事で最初に紹介するアルゴリズムは色々なところで Shor の離散対数問題を解くアルゴリズムとして載っているため、最初に参考にして書いたものである。にもかかわらず、途中に測定があるために Qiskit での実装が難しく、そのために測定を最後におこなうバージョンとして書いたものが次のアルゴリズム。しかしながら、ここまでの方法では実は求まらないという事実が判明し、結局 Shor の原論文に立ち向かうのが一番の近道だったと判明するのが後半である。要は、私の辿った道をともに歩んでほしいという文学的（？）なノリである。

背景（非推奨）

Shor のアルゴリズム における周期発見アルゴリズムの応用先の一つ。最速の古典アルゴリズムでは \(n\) （定義は次節の問題にて）に関する超多項式時間¹ (Shor の論文内には古典では \(\exp(O(\log p)^{1/3}(\log\log p)^{2/3})\) の時間計算量がかかるとある) かかるのに対し、 Peter Shor により \(\mathrm{poly}(n)\) 時間の計算量で解ける量子アルゴリズムが得られた。

問題（非推奨）

\(p\) を \(n\) ビットの素数とし²、 \(Z_p^\ast=\{1,\dots,p-1\}\) とする。 \(\alpha\) を \(Z_p^\ast\) の生成子とし³、 \(\beta\in Z_p^\ast\) とする。 \(p\), \(\alpha\), \(\beta\) が与えられたとき、 \(\alpha^d=\beta\) を満たす \(d\in[0, p-1]\) を求める⁴。

アルゴリズム（非推奨）

内容は基本的に Shor による周期発見アルゴリズムと同じである。ただし、二変数関数でそれぞれに周期が存在するというところが異なる。今回こそ [P. Shor, 1996] (p.20) に従いたかったが、よくわからなかったので、楕円曲線への適用をした [J. Proos and C. Zalka, 2004] のレビューにあった内容 (2章) に従う。このアルゴリズムは 2 つの量子アルゴリズムから構成される。

1 つ目の量子アルゴリズム: \(\alpha^q=1\) を満たす \(q\) を探す（非推奨）

本題の \(d\) を見つける前に、まずはじめに \(\alpha^q=1\) を満たす \(q\) を見つける必要がある。すなわち、 \(\alpha^x=\alpha^{x+q}\) という周期関数であるため、前回のショアの Factoring アルゴリズムにより得ることができる。

2 つ目の量子アルゴリズム: \(\alpha^d=\beta\) を満たす \(d\) を探す（非推奨）

1 つ目の量子アルゴリズムにより得られた \(q\) を考えると、 \(0\leq d\leq q\) となる。そのうえで \(f(x,y)=\alpha^x\beta^y\) を整数 \(x\), \(y\) について定義する。この関数 \(f\) は 2 つの独立な周期をもつ⁵。

$$f(x+q,y)=f(x,y),\ f(x+d,y-1)=f(x,y)$$

この関数 \(f\) を用いて、重ね合わせ状態

$$\frac{1}{q}\sum_{x=0}^{q-1}\sum_{y=0}^{q-1}\ket{x}\ket{y}\ket{0}$$

から

$$\frac{1}{q}\sum_{x=0}^{q-1}\sum_{y=0}^{q-1}\ket{x}\ket{y}\ket{\alpha^x\beta^y}\tag{1}$$

を得る。ここで、最後のレジスタ \(\ket{\alpha^x\beta^y}\) に測定をし、 \(\alpha^{x_0}\) が得られたとする。すると、

$$\alpha^x\beta^y=\alpha^x(\alpha^d)^y=\alpha^{x_0}$$

より、 \(x+dy\equiv x_0\pmod q\) 、つまり、 \(x=(x_0-dy)\pmod q\) となる。したがって、状態は

$$\frac{1}{\sqrt{q}}\sum_{y=0}^{q-1}\ket{x_0-dy}\ket{y}$$

となる。ここに量子フーリエ変換を適用し、

$$ \begin{align} &\frac{1}{\sqrt{q}}\frac{1}{q}\sum_{x',y'=0}^{q-1}\sum_{y=0}^{q-1}e^{2\pi i((x_0-dy)x'+yy')/q}\ket{x'}\ket{y'}\tag{2}\\ =&\frac{1}{\sqrt{q}}\sum_{x'=0}^{q-1}e^{2\pi ix_0 x'/q}\ket{x'}\ket{y'\equiv dx'\pmod q}\tag{3} \end{align} $$

を得る。ここで、式(2)から式(3)への変形では

$$ \sum_{a=0}^{n-1}e^{\frac{2\pi i}{n}ab}=\begin{cases} n & (b=b'n) \\ 0 & (\mathit{otherwise}) \end{cases} $$

を用いている（\(a\), \(b\), \(b'\), \(n\) は整数）。つまり、 \(dx'-y'\) が \(q\) で割り切れる場合にのみ \(y\) についての和の結果は \(q\) となり、そのときの \(y'\) の値は \(y'\equiv dx'\pmod q\) を満たす値でなくてはならない。そしてその値が 2 番目のレジスタに入っていることになる。

これを測定し、得られた \(x'\neq 0\), \(y'\) を用いて

$$d=y'(x')^{-1}\bmod q$$

と求まる。・・・のだが、次節で測定を最後に移動した場合⁶を書く。

アルゴリズム（測定を最後にした場合）（非推奨）

次に実装例を書こうと思ったら Qiskit で途中で測定する方法がよくわからなかったので、測定を最後に持っていった場合の 2 つ目（というか 1 つ目はもはや不要）の量子アルゴリズムを記す⁷。要は Factoring のときと同じ感じになる。

式(1)から始めるが、どうせ量子回路組むときはビット数のべき乗 \(N\) を分母にすることになるので、もうそのように書くと

$$\frac{1}{N}\sum_{x=0}^{N-1}\sum_{y=0}^{N-1}\ket{x}\ket{y}\ket{\alpha^x\beta^y}$$

量子フーリエ変換を適用し、

$$\frac{1}{N^2}\sum_{x'=0}^{N-1}\sum_{y'=0}^{N-1}\sum_{x=0}^{N-1}\sum_{y=0}^{N-1}e^{2\pi i(xx'+yy')/N}\ket{x'}\ket{y'}\ket{\alpha^x\beta^y}$$

3 つのレジスタ \(\ket{x'}\), \(\ket{y'}\), \(\ket{\alpha^{x_0}\beta^{y_0}}\) が測定される確率は

$$ \begin{align} &\frac{1}{N^4}\left|\sum_{x,y\ \text{where}\ \alpha^x\beta^y=\alpha^{x_0}\beta^{y_0}}e^{2\pi i(xx'+yy')/N}\right|^2\tag{4}\\ =&\frac{1}{N^4}\left|\sum_k\sum_\ell e^{2\pi i\{(x+kq+\ell d)x'+(y-\ell)y'\}/N}\right|^2\tag{5} \end{align} $$

ここで、式(4)から式(5)への変形には上で述べた 2 つの独立な周期を用いている。 Factoring のときと同様に測定で得られるのは \(\frac{qx'}{N}\) が整数かつ \(\frac{dx'-y'}{N}\) が整数を満たすものがほとんどなので、特に後者の条件より

$$dx'\equiv y'\pmod N$$

\(\gcd(x',N)=1\) のとき

$$d=y'(x')^{-1}\pmod N\tag{6}$$

より \(d\) は求まる。式(6)の解き方としては、拡張ユークリッド互除法より

$$x'd_0+Nd_1=1$$

を満たす整数 \(d_0,d_1\) を求め、その両辺に \(y'\) を掛けることで

$$x'd_0y'+Nd_1y'=y'$$

より

$$x'd_0y'\equiv y'\pmod N$$

であるから、 \(d=d_0y'\) と求まる。

これでは不完全なことがわかる例（非推奨）

\(p=7\) として \(Z_7^\ast=\{1,2,3,4,5,6\}\) とする。生成子 \(\alpha=3\), \(\beta=6\) とすると、答えは \(\alpha^3\equiv\beta\pmod 7\) より、 \(d=3\) となる。これを shor のアルゴリズムにより求める。

今回のアルゴリズム（Qiskit プログラムは後述）を動かすことにより以下のプロットを得る。

（やはり若干読みづらいが、例えば右端 100111 は \(\ket{x'}\ket{y'}=\ket{111}\ket{100}=\ket{7}\ket{4}\) を表す）

最も大きいピークが 2 本、次点で大きいピークが 4 本存在する。そう、これも前回と同じである。ゆえに、やはり \(u=\frac{qx'}{N},v=\frac{dx'-y'}{N}\) について表を書くと次のようになる。この表は簡単に

def f(u,v):
    x = u * 8 / 6
    y = 3 * x - 8 * v
    return x,y

for u in range(20):
    for v in range(20):
        x,y = f(u,v)
        if 0<=x<8 and 0<=y<8:
            print(f'(u,v)=({u},{v}): (x,y)=({x},{y})')

で求めた（\(3^6\equiv 1\pmod 7\) より \(q=6\)）。

\((u,v)\)	\(x'\)	\(y'\)
(0,0)	0	0
(1,0)	1.333	4
(2,1)	2.666	0
(3,1)	4	4
(4,2)	5.333	0
(5,2)	6.666	4

より、やはり最も大きいピークは \(x',y'\) が整数のところであり、次点はそれに近い小数値となっている。ゆえに、測定で得られた \(x'\) を用いて拡張ユークリッド互除法をおこなっても正しい \(d\) を得ることは難しい。現に、今回のピークの値を用いてもうまく \(d\) を求めることはできない。

では、 \(N=8^2=64\) でおこなうとどうだろうか。再び後述する Qiskit プログラムでピークを描いてみると以下のようになる。

ここで得られた大きなピークを用いて上で述べたように拡張ユークリッド互除法をおこなっても、やはり残念ながら正しく \(d\) を得ることはできない。

測定で得られた整数値から \(d\) を計算（おそらく完全版）（非推奨）

仕方がないので、結局 Shor の原論文を参考にするのがこういうときは良いだろう。論文の p.21 を見ると

$$T=dx'+y'-\frac{d}{p-1}\{x'(p-1)\}_N$$

という式が出てくる（本記事に各記号は合わせた）。ここで、 \(\{z\}_N\) は \(z\pmod N\) のことである。

これは次の不等式を満たす。

$$|\{T\}_N|=|dx'+y'-\frac{d}{p-1}\{x'(p-1)\}_N-jN|\leq\frac{1}{2}$$

ここで、 \(j\) は \(T/N\) に最も近い整数である。ゆえに、

$$ \begin{align} -\frac{1}{2}&\leq dx'+y'-\frac{d}{p-1}\{x'(p-1)\}_N-jN \leq\frac{1}{2}\\ -\frac{1}{2N}&\leq \frac{y'}{N} + d\left(\frac{x'(p-1)-\{x'(p-1)\}_N}{(p-1)N}\right)-j\leq\frac{1}{2N}\\ -\frac{1}{2N}&\leq \frac{y'}{N} + d\left(\frac{x'(p-1)-\{x'(p-1)\}_N}{(p-1)N}\right)\leq\frac{1}{2N}\pmod 1 \end{align} $$

最後の変形では \(\pmod 1\) を用いた。これにより \(j\in\mathbb{Z}\) は \(1\) で割った余りが \(0\) なので消える。ここで、 \(\{x'(p-1)\}_N=x'(p-1)-c'N\) (\(c'\in\mathbb{Z}\)) なので

$$-\frac{1}{2N}\leq \frac{y'}{N} + \frac{dc'}{p-1}\leq\frac{1}{2N}\pmod 1$$

それから \(\frac{y'}{N}\) を連分数展開（など）により分母を \(p-1\) にする（このとき得られる分子を \(y''\in\mathbb{Z}\) とする）

$$-\frac{1}{2N}\leq \frac{y''+dc'}{p-1}\leq\frac{1}{2N}\pmod 1$$

ここから原論文では \(c'\) で割る \(\pmod {p-1}\) とあったのだが、（ちょっとなにいってるかわからなかったので）ここでは \(N\gg 1\) つまり、左辺と右辺がほぼ 0 であることと \(\pmod 1\) を活かして $$\frac{y''+dc'}{p-1}=m\in\mathbb{Z}$$ となるような \(d\) を求めるアプローチをとる。式変形して $$ \begin{align} y''+dc'&=m(p-1)\\ dc'-m(p-1)&=-y'' \end{align} $$ 最後の不定方程式を拡張ユークリッド互除法により解くことで \(d\) が求まる。

Qiskit による実装（非推奨）

ここまでの内容を Qiskit を用いて実装したのが以下になる。コードの全体は GitHub にある。

from typing import Optional
import math

import numpy as np
import matplotlib.pyplot as plt
from sympy import gcdex
from qiskit import QuantumCircuit, Aer, transpile, assemble
from qiskit.visualization import plot_histogram
from sympy import Rational
from sympy.ntheory.continued_fraction import continued_fraction, continued_fraction_convergents

from qft import qft
from elementary import ax_modM

def discrete_log(alpha: int, beta: int, p: int, N_len: Optional[int] = None, show_hist: Optional[bool] = True) -> int:
    # find d s.t. alpha^d = beta
    if N_len is None:
        N_len = int(np.ceil(np.log2(p ** 2)))
    N = 2 ** N_len
    qc = QuantumCircuit(11 * N_len - 2, N_len * 2)

    qc.h(range(N_len * 2))

    qc.append(ax_modM(a=alpha, M=p, N_len=N_len), list(range(N_len)) + list(range(N_len * 2, 11 * N_len - 2)))
    qc.append(ax_modM(a=beta, M=p, N_len=N_len, x_0_at_first=False), range(N_len, 11 * N_len - 2))

    qc.append(qft(n=N_len), range(N_len))
    qc.append(qft(n=N_len), range(N_len, N_len * 2))

    qc.measure(range(N_len * 2), range(N_len * 2))

    backend = Aer.get_backend('aer_simulator_matrix_product_state')
    qc = transpile(qc, backend)
    job = backend.run(qc, shots=10000)
    hist = job.result().get_counts()

    if show_hist:
        plot_histogram(hist)
        plt.show()

    for measured_key, _ in sorted(hist.items(), key=lambda x: x[1], reverse=True):
        x = int(measured_key[-N_len:], 2)
        y = int(measured_key[-N_len * 2:-N_len], 2)

        for fraction in continued_fraction_convergents(continued_fraction(Rational(y, N))):
            if fraction.denominator == p - 1:
                yy = fraction.numerator
            elif (p - 1) % fraction.denominator == 0:
                yy = fraction.numerator * ((p - 1) // fraction.denominator)
            else:
                continue

            c = Rational(x * (p - 1) - (x * (p - 1) % N)) / N
            if c == 0:
                continue

            d, m, _ = gcdex(c, 1 - p)
            d *= -yy
            m *= -yy
            if alpha ** d % p == beta:
                return d
            d += 1 - p
            if alpha ** d % p == beta:
                return d

    print('d was not found?!')

if __name__ == '__main__':
    print(discrete_log(alpha=3, beta=6, p=7, N_len=3))

これを実行すると、確かに \(d=3\) が得られる。