💯

【AtCoder】リストへ複数の要素を追加したり、取り出して合計を計算（配列処理の高速化）

2025/07/05に公開

問題

配列に要素を追加するクエリ1、取り出して計算するクエリ2を処理する
クエリ1：配列の末尾にxという要素をc個追加する
クエリ2：配列の先頭からk個の要素を取り出し、合計値を出力

解答コード

N = int(input())
A = []
B = []
for _ in range(N):
    q = [int(i) for i in input().split()]
    if q[0] == 1:
        A.append({'c':q[1], 'x':q[2]})
    else:
        B.append(q[1])
# 配列反転
A = A[::-1]

for k in B:
    ans = 0
    while k > 0:
        
        if k <= A[-1]['c']:
            ans += k * A[-1]['x']
            A[-1]['c'] -= k
            
            if A[-1]['c'] == 0:
                del A[-1]
                
            break
        elif k > A[-1]['c']:
            k -= A[-1]['c']
            ans += A[-1]['c'] * A[-1]['x']
            del A[-1]
        
    print(ans)

攻略ポイント

配列に要素を追加、取り出して出力という簡単な操作で、一見すると初心者でも実装可能なように思える。しかし入力例3の時点でかなり重い処理になっており、配列処理を効率的に行う必要がある。

解説

解説1：クエリを2種類を分ける

クエリを1回ごとに処理せず、まずはA,Bという配列に振り分ける。

A = []
B = []
for _ in range(N):
    q = [int(i) for i in input().split()]
    if q[0] == 1:
        A.append({'c':q[1], 'x':q[2]})
    else:
        B.append(q[1])

解説2：クエリの通りに処理しない

ほぼクエリと同じ形で配列Aに突っ込む。

A.append({'c':q[1], 'x':q[2]})

追加クエリは配列末尾にxという要素をc個追加するというもの。実際にcの数だけ配列へ追加しているとcが膨大になった時、実行時間制限内に処理しきれない。（削除クエリも同じ）
実際に配列に追加しなくても、 xがc個ある と分かっているのだから、この形のままでも計算できる。

解説3：配列の取り出しを末尾からやる

クエリ2は「配列の先頭からk個の要素を取り出し、合計値を出力」。pythonでは 先頭からの取り出しはインデックスの振り直しが発生して低速になる。 そのため、配列を反転させて末尾から取り出せるようにする。

A = A[::-1]

解説4：A配列の末尾から必要な分だけ取り出して削除

配列の末尾A[-1]のcの数に着目。kとの大小関係によって3パターン発生する。

kとcが同数：取り出し終了、末尾を削除
cの方が大きい：取り出し終了、cをkの分だけ減らす
kの方が大きい：取り出し続行、末尾を削除、cをkの分だけ減らす

if k <= A[-1]['c']:
    ans += k * A[-1]['x']
    A[-1]['c'] -= k
    
    if A[-1]['c'] == 0:
        del A[-1]
        
    break
elif k > A[-1]['c']:
    k -= A[-1]['c']
    ans += A[-1]['c'] * A[-1]['x']
    del A[-1]

問題