AtCoder ABC 140 B – Buffet Python解説

スポンサーリンク

Buffet

高橋くんは N 種類の料理が食べ放題のビュッフェに行き、全種類の料理 (料理 1, 料理 2, 料理 N) を 1 度ずつ食べました。
高橋くんが i  番目に食べた料理は料理 Ai でした。
高橋くんは、料理 i (1≤iN) を食べると満足度 Bi を得ます。
また、料理 i (1≤iN−1) を食べた直後に料理 i+1 を食べると満足度 Ci を追加で得ます。
高橋くんが得た満足度の合計を求めてください。

AtCoder Beginner Contest 「Buffet」

料理を食べると満足度が上がり、食べた直後に料理i+1を食べると満足度が追加されます。ですので、ひとつ前に食べた料理の番号を保持する変数を用意し、追加の満足度に気を付けてコードを書きます。

n = int(input())

a = list(map(int, input().split()))
b = list(map(int, input().split()))
c = list(map(int, input().split()))

ans = 0
order = -2

for i in a:
    ans += b[i - 1]
    if order + 1 == i:
        ans += c[order - 1]
        order = i 
    else:
        order = i 
        
print(ans)

ansに満足度、orderに直前に食べた料理の番号を入れます。for文の中では食べた料理の満足度をansに足していき、ひとつ前に食べた料理の番号+1がiと同じになったら追加の満足度c[order – 1]をansに加えていきます。
そして、ひとつ前に食べた料理の番号をiで更新すればコードの完成となります。