AtCoder-ABC216 E - Amusement Park【Python解答例】 - ebisukeプログラミング初心者脱出黙示録

AtCoder Beginner Contest216のE問題についてPythonの解答例を記事にしていきます。
AtCoder Beginner Contest 216 - AtCoder

AtCoder Beginner Contest216 E - Amusement Park

問題文

髙橋君は遊園地に遊びに行きました。
この遊園地には N 個のアトラクションがあり、i 個目のアトラクションの「楽しさ」の初期値は A i です。

髙橋君が i 個目のアトラクションに乗ると、以下の現象が順番に起きます。

髙橋君の「満足度」に、i 個目のアトラクションの現在の「楽しさ」が加算される。
i 個目のアトラクションの「楽しさ」が、1 減少する。
髙橋君の「満足度」の初期値は 0 です。髙橋君はアトラクションに合計 K 回まで乗ることができます。
最終的な髙橋君の「満足度」の最大値はいくつですか？

なお、髙橋君の「満足度」はアトラクションに乗ること以外で変化しません。

制約

・1≤N≤10 ^5
・1≤K≤2×10 ^9
・1≤A i ≤2×10 ^9
・入力は全て整数

解答例

N, K = map(int,input().split())

A = list(map(int, input().split()))

# 降順で最後の要素が0
A.sort(reverse=True)
A.append(0)

# 搭乗回数
k = 0
# 合計満足度
score = 0
# 搭乗回数が余るかどうか
flag = True

for i in range(N):
    # 次の要素と一致するまでに必要な搭乗回数
    tmp = (i + 1) * (A[i] - A[i+1])
    # tmpがKより大きければberak
    if k + tmp > K:
        break
    # スコアとkを更新
    score += (i + 1) * (A[i]*(A[i]+1) - (A[i+1])*(A[i+1]+1)) // 2
    k += tmp
else:
    flag = False

Kが0になるまでのスコア計算
if flag:
    div = (K-k) // (i+1)
    mod = (K-k) % (i+1)
    tmp = A[i]
    
    score += (i+1) * (tmp*(tmp+1) - (tmp-div)*(tmp-div+1))//2
    score += mod * (tmp-div)

print(score)