yukicoder No.16 累乗の加算

※noteに記事を投稿していましたが，数式を書きたくなることがたまにあるのでHatena Blogも試してみることにしました。

note.com

下記問題に挑戦した。問題詳細は以下。

yukicoder.me

愚直に実装したらTLEでした（そりゃそうだ）。

x, N = map(int,input().split())
A = list(map(int,input().split()))
mod = 1000003
ans = 0
for a in A:
    ans += (x**a)%mod
print(ans)

ということで，解説記事を読むと，繰り返し自乗法を使いなさいとのこと。少し調べると手元の本にも載っていました。

• プログラミングコンテスト攻略のためのアルゴリズムとデータ構造 （p.445 18.3 べき乗）

• はじめての数論　原著第3版 　（第16章）

後者の本をとりあえず軽く読んで理解した（つもり）。要は

1. のリストをに対して生成する
2. 指数を二進展開する ということを行えばよい。 1.の計算が直前の数を二乗（自分自身との積なので自乗とも）することを繰り返すことからこういう名前なのだと思われる。

愚直に実装した。

#aのべき乗のリストの生成（繰り返し二乗法）
def mkPow(a,k):
    P[0] = 1
    i = 1
    while i <= k:
        P[i] = a
        a = (a**2)%mod
        i *=2

#nの2進展開に含まれる２のべき乗のリストを返す
def exTwo(n):
    ex = []
    i = len(P)
    kList = list(P.keys())
    for i in reversed(kList):
        if  n - i >= 0:
            n -= i
            ex.append(i)
        if n == 0:
            return ex

mod = 1000003
x, N = map(int,input().split())
A = list(map(int,input().split()))
P = {}
mkPow(x,max(A)+1)
ans = 0
for a in A:
    lis = exTwo(a)
    xPow = 1
    for l in lis:
        xPow *= P[l]
        xPow %= mod
    ans += xPow
    ans %= mod
print(ans%mod)

べき乗のリストをどこまで作ればよいか少し悩んだが，pythonのdictを使うことにした（指数部とべき乗の値をセットで保持しやすい）。 無事AC。 そのあと，少し調べるとpythonにはpowという組み込み関数があるので，これを使って書き直してみた。

mod = 1000003
x, N = map(int,input().split())
A = list(map(int,input().split()))
ans = 0
for a in A:
    ans += pow(x,a,mod)
print(ans%mod)

コードめちゃ短い。もちろんＡＣ。標準の関数って優秀ですね～。