スタックとして使うときのlist.extendの速度 - Python

何が原因でこうなるのか、あまりつかみ切れていないが、とりあえず発生することは事実。

以下のコード、単に配列Aに配列Bの3の倍数のみを付与している。

実行時間は、extendの方が若干速い。何となく、リスト内包表記が速いことからも頷ける。

def append(A, B):
    for b in B:
        if b % 3 == 0:
            A.append(b)


def extend(A, B):
    A.extend(b for b in B if b % 3 == 0)


A = [0] * 1000
B = list(range(10 ** 6))
t1 = timeit('append(A, B)', number=100, globals=globals())
t2 = timeit('extend(A, B)', number=100, globals=globals())
print(t1)  # => 9.2940 
print(t2)  # => 8.5101

（ただ、今回は3の倍数としたが、条件に当てはまる要素が多いほどどちらも時間がかかるようになる）

では、こちらはどうか。

木構造の深さ優先探索を行っているコードである。スタックに隣接頂点を追加しながら探索を深めていっている。

頂点は (コスト, 頂点ID, 親頂点ID) で情報を持ち、親頂点には戻らないようにしている。

結果は、extendを用いた方が、2倍以上遅い。先ほどと似たようなコードなのに、何が違うんだろう。。。

def dfs1(links, s):
    stack = [(0, s, -1)]
    while stack:
        cost, v, p = stack.pop()
        cost += 1
        stack.extend((cost, u, v) for u in links[v] if u != p)


def dfs2(links, s):
    stack = [(0, s, -1)]
    while stack:
        cost, v, p = stack.pop()
        cost += 1
        for u in links[v]:
            if u != p:
                stack.append((cost, u, v))


n = 200000
links = [set() for _ in range(n)]
for i in range(n - 1):
    links[i].add(i + 1)
    links[i + 1].add(i)

t1 = timeit('bfs1(links, 0)', globals=globals(), number=10)
t2 = timeit('bfs2(links, 0)', globals=globals(), number=10)
print(t1)  # => 1.1736078369431198
print(t2)  # => 0.49144828389398754

他に調べたこととして、

• extend()の引数をgeneratorではなく、[ ] で括ってリストにした方が、幾分か速くなる（が、まだappendの方が速い）
• 親子情報が既知で、if u != p で除かなくてよい場合でも、やはりextendの方が遅い
• cost情報も取り払って「stack.extend(u for u in links[v])」にしても、あまり変わらず
• 「stack.extend(links[v])」まですると、やっとappendより高速になるにはなった
  • が、ここまですると何の情報も無くなり、探索的には意味をなさなくなる

原因が一目には分からないが、とりあえずこのようなstack的な使い方をする場合は、extendは避けた方がよさそう。

幅優先探索を行う際は List の代わりに collections.deque を用いて探索順を先入れ先出しにしたりするが、これも同様、extendの方が遅い結果となる。

• 20万頂点の直線形の木を、端から端までBFSする時間（10回の平均）
  • extend: 0.1879s
  • append: 0.0843s
• 20万頂点のスターグラフを、中心からBFSする時間（10回の平均）
  • extend: 0.1688s
  • append: 0.0818s

実験コード

from collections import deque
from timeit import timeit


def bfs1(links, s):
    q = deque([(0, s, -1)])
    while q:
        cost, v, p = q.popleft()
        cost += 1
        q.extend((cost, u, v) for u in links[v] if u != p)


def bfs2(links, s):
    q = deque([(0, s, -1)])
    while q:
        cost, v, p = q.popleft()
        cost += 1
        for u in links[v]:
            if u != p:
                q.append((cost, u, v))


n = 200000
links = [set() for _ in range(n)]
for i in range(n - 1):
    links[i].add(i + 1)
    links[i + 1].add(i)

t1 = timeit('bfs1(links, 0)', globals=globals(), number=10)
t2 = timeit('bfs2(links, 0)', globals=globals(), number=10)
print(t1)  # => 1.8790318
print(t2)  # => 0.8429923