(ABC368)C - Triple Attack と pop vs pointerの速度の話

はじめに

ABC368のC問題です。
またしてもTLE（実行時間制限超過）が発生させて
ハマった話です。

自分の回答（1回目）

ABC368_C_1.py

N = int(input())
E = list(map(int,input().split()))
T = 0
while E:
  while E[0] > 0:
    if E[0] >= 5:
      T += 3
      E[0] -= 5
    else:
      T += 1
      if T % 3 == 0:
        E[0] -= 3
      else:
        E[0] -= 1
  E.pop(0)

print(T)

これは想定内です。
入力値が
$1≤N≤2×10^5$
$1≤H≤10^9$
ですから。計算量は $O(10^{14})$ であろう。

改善その１

すこし考えます。
$T$を3つセットで考えると、$(1,1,3)$、$(1,3,1)$、$(3,1,1)$のいずれの並びも$H$の体力を$5$減らせられるので、$H$が$5$以上の場合は、まとめて考えることができる。

ABC368_C_2.py

N = int(input())
E = list(map(int,input().split()))
T = 0
while E:
  while E[0] > 0:
    if E[0] >= 5:
      T += E[0] // 5 * 3
      E[0] = E[0] % 5
    else:
      T += 1
      if T % 3 == 0:
        E[0] -= 3
      else:
        E[0] -= 1
  E.pop(0)

print(T)

これはショックでした。
最大で計算量を$1/5$にしたと思ったのに。

ここでしばらく思考停止。

改善その２

ここでふとpopか？
配列の削除って処理遅いのかもと思い、popをやめてみる。

ABC368_C_3.py

N = int(input())
E = list(map(int,input().split()))
T = 0
for i in range(len(E)):
  while E[i] > 0:
    if E[i] >= 5:
      T += E[i] // 5 * 3
      E[i] = E[i] % 5
    else:
      T += 1
      if T % 3 == 0:
        E[i] -= 3
      else:
        E[i] -= 1

print(T)

無事通過です。

追加検証（popとポインタ移動の速度比較）

生成AIと一緒に検証してみる。

popとpointerの速度検証

import timeit
import random

def test_pop_vs_pointer(list_size):
    test_list = [random.randint(1, 100) for _ in range(list_size)]
    
    # pop(0)の速度を測定
    pop_time = timeit.timeit(
        lambda: test_list.pop(0) if test_list else None,
        number=list_size
    )
    
    # ポインタを動かす速度を測定
    pointer_time = timeit.timeit(
        lambda: [test_list[i] for i in range(min(list_size, len(test_list)))],
        number=1
    )
    
    return pop_time, pointer_time

# テストするリストのサイズ
list_sizes = [10, 100, 1000, 10000, 100000]

print("リストサイズ | pop(0)の時間 | ポインタ移動時間 | 差")
print("-" * 55)

for size in list_sizes:
    pop_time, pointer_time = test_pop_vs_pointer(size)
    time_difference = pop_time - pointer_time
    print(f"{size:11d} | {pop_time:10.6f} | {pointer_time:10.6f} | {time_difference:10.6f}")

結果

リストサイズ	pop(0)の時間	ポインタ移動時間	差
10	0.000006	0.000005	0.000001
100	0.000025	0.000004	0.000021
1000	0.000313	0.000006	0.000307
10000	0.009507	0.000006	0.009501
100000	0.949552	0.000010	0.949542

一目瞭然ですね。

おわりに

計算の部分に注目しすぎてハマってしまったという話でした。

Pythonの知識が１つ増えました！