AtCoder Beginner Contest 141のDをPython3で解いたメモ

問題↓
D - Powerful Discount Tickets

とりあえず自分が書いたやつ(TLE)

最安にしたいんだから, 現状最大のやつを値段半分にしていけばいいかと考えて書いた.
悪くはないと思うんだけどTLE. 時間切れでコンテスト中はここまで.

N, M = map(int, input().split(" "))
A = [int(i) for i in input().split(" ")]

for i in range(M):
    tmp_index = A.index(max(A))
    A[tmp_index] = A[tmp_index] / 2

import math
B = [math.floor(A[i]) for i in range(len(A))]
print(sum(B))

heapqモジュールを使えば最大値の取り出しが高速になるらしい

Pythonでヒープ構造を使う（heapqモジュール）

heapqは最小値を高速で見つけるものらしいので, とりあえずリスト内の値すべてに-1をかけておくというのが常套手段っぽい。
heapq.heapify(list)でソート, heapq.heappop(リスト)で最小値の値を取得&リストからその値を除去, heapqpush(リスト, 値）でリストに値を追加とのこと.

import math
import heapq

N, M = map(int, input().split(" "))
A = [-int(_) for _ in input().split(" ")]

for i in range(M):
    heapq.heapify(A)
    heapq.heappush(A, heapq.heappop(A)/2)

B = [math.floor(-A[i]) for i in range(len(A))]
print(sum(B))

しかし結果はほぼ変わらず(TLE)

math関数が足を引っ張っている可能性

切り捨てを行うのにmath.floorでやってたけれども遅いらしい.
Python で割り算をするときの
切り上げと切り捨て

これが一番高速らしい

# 切り捨て
4 // 3

# 切り上げ
-(-4 // 3)

というわけでそっちを修正したバージョンを

N, M = map(int, input().split(" "))
A = [int(_) for _ in input().split(" ")]

for i in range(M):
    tmp_index = A.index(max(A))
    A[tmp_index] = A[tmp_index] // 2    # ここが修正箇所

print(sum(A))

結果はTLE.

heapqの書き方が悪いだけだった

毎回, heapq.heapifyやってるせいで計算時間がかかりすぎてるっぽいので修正

import heapq

N, M = map(int, input().split(" "))
A = [-int(_) for _ in input().split(" ")]

heapq.heapify(A)

for i in range(M):
    heapq.heappush(A, -(-heapq.heappop(A)// 2))

print(-sum(A))

これでAC.