競プロ初中級者100問

Square869120Countest#1E - 散歩 (E869120 and Path Length)

この記事で使うアルゴリズム

累積和

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

https://atcoder.jp/contests/s8pc-1/tasks/s8pc_1_e

この問題では巨大な数の冪剰余の計算が求められます。

方針

  1. 冪剰余を求める
    • pow(m,n)%MODのnが巨大である場合、高速に計算する方法があります。こちらの記事で解説しています。
    • Pythonでは標準で実装されおり、pow(m,n,MOD)とすれば良いです。
  2. 累積和を求める
    • 仮想的な街0を考え、街0から任意の街iまでの距離を事前に計算しておきます（acc[i-1]とする）。そうすることで、任意の街(i,j)の間の距離をabs(acc[i-1]-acc[j-1])のように計算量O(1)で求められます。
    • 標準モジュールのitertools.accumulateが適当でしょう。

コード

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from itertools import accumulate
MOD = 10**9+7
N, Q = map(int, input().split())
A = list(map(int, input().split()))
C = list(map(int, input().split()))

D = [0]
for n in range(N-1):
    d = pow(A[n],A[n+1],MOD) # pow(A[n],A[n+1])%MOD より高速
    D.append(d)
acc = list(accumulate(D)) # 累積和

C.insert(0,1) # 始点
C.append(1) # 終点
ans = 0
for q in range(len(C)-1):
    ans += abs(acc[C[q+1]-1]-acc[C[q]-1]) % MOD
print (ans%MOD) # 最後のMODを忘れずに

記事情報

  • 投稿日：2020年4月30日
  • 最終更新日：2020年5月8日