やってみたらなんとかなる

プログラミングをする上で調べたこととかやったこととか

Dequeを使うといいのかも -AtCoder学習日記 #5

概要

今回はABC216-Dについてです。 今回はうまくいったと思ってたんですけどね...やはりTLEに悩まされます。

ACにならなかったコード

n, m = map(int, input().split())

nodes = [set() for _ in range(n+1)]

for _ in range(m):
  k = int(input())
  pipes = list(map(int, input().split()))
  for i in range(k):
    for j in range(i+1, k):
      nodes[pipes[i]].add(pipes[j])
      if pipes[i] in nodes[pipes[j]]:
        print("No")
        exit()

print("Yes")

筒の中で出てくる色の順番が異なる部分が出てきたらダメだと思って組んだんですけどTLEになってしまいました...

解説の解説

解説を私なりに要約すると「書いてあることをそのままやればいいよ。ただ計算量(実行時間)に気をつけてね」ということみたいです。
なので今回は前から気になっていたDeque型を使ってみたいと思います。

Dequeについて

Deque とは、スタックとキューを一般化したものです (この名前は「デック」と発音され、これは「double-ended queue」の省略形です)。

ちなみに、Listでも同じことが可能です。しかし、Dequeを用いた方が圧倒的に早い模様です。それぞれの計算量は以下のようになってます。
ListとDequeで計算量が変わるところは赤くしています。

操作 List Deque
先頭に追加 O(n) O(1)
最後尾に追加 O(1) O(1)
先頭を取り出す O(n) O(1)
最後尾を取り出す O(1) O(1)

先頭についての操作はDequeの方が上手みたいですね。

ACなコード

from collections import deque

n, m = map(int, input().split())
pipes = []
for _ in range(m):
  k = input()
  pipes.append(deque(map(int, input().split())))

queue = deque([(q.popleft(), i) for i, q in enumerate(pipes)])
# queueに各筒の先頭の球を入れる。

balls = [-1] * (n + 1)
# 入ってた筒と球が2つ出たかどうかの判定

while queue:
  
  a, i = queue.popleft()

  if balls[a] == -1:
    balls[a] = i
    # 1つ目の球なら球の入っていた筒の番号を保存

  else:
    j = balls[a]
    balls[a] = -2
    # 2つ目の球なら2つ出たことを記録して
    # jに1つ目の球が入っていた筒の番号
    # iに2つ目の球が入っていた筒の番号を記録

    if pipes[j]: 
      queue.append((pipes[j].popleft(), j))
      #もしj番目の筒に球があるならqueueに入れる
      
    if pipes[i]: 
      queue.append((pipes[i].popleft(), i))
      #もしi番目の筒に球があるならqueueに入れる

if sum(balls) == -2 * n - 1:
  print('Yes')
else:
  print('No')

今回からコードの説明をコメントで書き込んでみました。少しはわかりやすくなっている...はず...です。

まとめ

  • 先頭と最後尾を扱う配列の場合はListよりもDequeの方が計算量が少ない!
  • 問題によっては問題文通りに進めるのも良いこともある

参考文献

docs.python.org

wiki.python.org