はむ吉（のんびり）の練習ノート

プログラミングとことばに関する話題を中心に，思いついたこと，試してみたこと，学んだことを，覚え書きを兼ねてまとめます．その際に役立った，技術書，参考書，辞書，機器などの紹介も行います．

square869120Contest #1 E. 散歩：解き方とPython 3による解答例

プログラミング 競技プログラミング AtCoder

1月24日の晩に、AtCoder上でプログラミングコンテストsquare869120Contest #1が開かれました。本記事では、このコンテストで出題された8つの問題のうち、E: 散歩 (E869120 and Path Length) - square869120Contest #1 | AtCoderを取り上げ、私の頭の整理を兼ねて、その解法およびPython 3による解答例を示します。

問題文の解釈

問題文では、街の番号が1始まりとなっています。ところが、配列の添字は通常0から始まります。そこで、入力は適宜変換するとして、街の番号を0始まりに改めることにします。すると、この問題は以下のように置き換えられます。

街がN個あり、各街i\:(i=0,1,\cdots,N-1)に整数値a_iが結びつけられている。また、2つの街j, j+1\:(j = 0, 1, \cdots, N-2)の間には長さw_j={a_j}^{a_{j+1}}の道がある。いま、街0を出発し、Q個の街c'_1, c'_2, \cdots, c'_k, \cdots, c'_Qを順に訪問して街0に戻るときに移動する道のりをdとする。また、M=10^9 + 7とする。入力としてN, Q, \{a_i\}, \{c_k\}（ただしc_k = c'_k + 1）が与えられるので、dをMで除したときのあまりd'を出力せよ。

方針

街sから街tまでの道のりは、S_{s, t} := \sum ^{t-1} _{j=s} w_jで与えられます*1。いま、w_jは非常に大きな値をとるおそれがあるので、w'_j = w_j \% Mに対してS'_{s, t} := \sum ^{t-1} _{j=s} w'_jを定めます*2。すると、合同式の性質に注意すれば、これを\{c'_k\}に従って足し合わせ、その結果をMで割ったあまりを求めると、それがd'となることがわかります。

使用する技法

上で定めたw'_jは冪剰余の形をしており、これを高速に計算する方法があります。また、上の方法では部分配列の総和S'_{s, t}を繰り返し計算しますが、これを効率よく行う技法として、累積和があります。

冪剰余

一般に、b^e \% m \: (b,e,m\in\mathbb{Z}; b,m>0)を、mを法としたbのe冪剰余と呼びます*3。wikipedia:冪剰余には、これを高速に計算するアルゴリズムに関する記述があります。とはいえ、Pythonのpow関数を用いれば、これらのアルゴリズムを実装しなくても高速に冪剰余の計算を行えます。

累積和

配列\{x_n\}（長さN'）の第p項から第q項まで（ただしp,q\in\mathbb{Z}; 0\le p\le q）の総和T_{p, q} := \sum ^{q} _{i=p} x_iをさまざまなp,qに対して何度も求める際には、累積和の方法が有効です。漸化式y_0 = 0,\,y_{m + 1} = y_m + x_mにより累積和の配列\{y_m\}を作成すれば（計算量O(N')）、T_{p,q} = y_{q + 1} - y_pにより上記の総和を求めることができます（計算量O(1)）*4。

Python 3による解答例

以上の方針および技法に基づき、私が作成したコードは以下の通りです*5。

コード作成にあたり注意した点は以下の通りです。

  • 配列をarrayオブジェクトに格納し、（主にメモリ面での）効率の向上を図る
  • 可読性を損なわない限り内包表記やlambdaを利用して、for文の使用を控え、処理速度の向上を図る

感想

冪剰余および累積和については、Project EulerやAizu Online Judgeの問題を通じて知っていました。そのため、コンテスト当時、途中で何度かWrong Answerを発生させたものの、最終的にAcceptedを得ることができました。普段、私はプログラミングコンテストで、アルゴリズムの知識をあまり要求しない簡単な問題しか解けずに終わることが多いので、自ら冪剰余と累積和を思い出し、この問題を解けたのは喜ばしいことでした。

*1:いくつかの小さなs, tについて確認すればわかりやすいでしょう。

*2:ここで、 p \% qはpをqで割ったあまりとします。

*3:wikipedia:冪剰余によります。

*4:累積和の使い方を学ぶ – himajinworks::blogによります。

*5:私がコンテスト開催中に提出したコードを修正したものです。