Pythonで大きな配列を高速に処理したい

はじめに

Pythonで競技プログラミングにチャレンジしています。（競技プログラミング歴半年くらい）
競技プログラミングでは配列をよく使うのですが、配列サイズが大きくなるに連れて処理時間がかかるようになり、速度で不利なPythonではかなり苦労しています。

配列に関する処理をどうやって速く実装できるかいろいろ調べたり実験したりした結果を自分のためにも残しておこうと思います。

環境

MacBook Pro (15-inch, 2017)
macOS 10.14.1
Python 3.7.1

配列の実装方法

N×Nの2次元配列での実装です。
（競技プログラミングでの想定問題は、N個の頂点／M個のパスのグラフ問題のイメージ）

リスト（list）

素直に実装すればこうかな？

list.py

graph = [[0 for _ in range(N)] for _ in range(N)]

NumPy

効率的な数値計算を行うための型付きの多次元配列をサポートした高速なライブラリ

numpy.py

import numpy as np
graph = np.zeros((N,N),dtype=np.int)

ディクショナリ【シングル】（defaultdict）

デフォルト値定義可能なディクショナリ型。該当するキーがなくてもエラーを返さず定義したデフォルト値を返す便利なやつ。
ディクショナリをリストに保持して２次元配列にする感じ。

defaultdict.py

from collections import defaultdict
graph = [defaultdict(int) for _ in range(N)]

ディクショナリ【ダブル】（defaultdict）

ディクショナリを入れ子にした感じ。

defaultdict2.py

from collections import defaultdict
graph = defaultdict(lambda: defaultdict(int))

実測

N×Nの2次元配列で、M×2回のランダムアクセスをした場合の速度を比較しました。

1000x1000配列,1000x2回ランダムアクセス

リストオンリー
初期化：0.039746秒
ランダムアクセス：0.000465秒

numpy
初期化：0.000017秒
ランダムアクセス：0.003246秒

シングルディクショナリ
初期化：0.000341秒
ランダムアクセス：0.000331秒

ダブルディクショナリ
初期化：0.000002秒
ランダムアクセス：0.000673秒

↑リストの初期化に時間がかかってる感じ。内包表記は速いって聞いてたけどそうでもないかも？
シングルディクショナリの初期化でもリスト部分が足を引っ張っているみたい。
numpyは初期化が速い。ダブルディクショナリはもっと速い。

1000x1000配列,1000000x2回ランダムアクセス

リストオンリー
初期化：0.039718秒
ランダムアクセス：0.654044秒

numpy
初期化：0.000023秒
ランダムアクセス：0.441909秒

シングルディクショナリ
初期化：0.000417秒
ランダムアクセス：0.821422秒

ダブルディクショナリ
初期化：0.000006秒
ランダムアクセス：0.872100秒

↑配列サイズは同じで、アクセス回数を増やした感じ。
ランダムアクセスではnumpyが一番速く見える。要素へのアクセスは遅いって聞いた気がするけど、そこまででもないのかな？

10000x10000配列,10000x2回ランダムアクセス

リストオンリー
初期化：3.741292秒
ランダムアクセス：0.006884秒

numpy
初期化：0.000020秒
ランダムアクセス：0.052834秒

シングルディクショナリ
初期化：0.004078秒
ランダムアクセス：0.005256秒

ダブルディクショナリ
初期化：0.000004秒
ランダムアクセス：0.008571秒

↑配列サイズが大きくなると、リストの初期化がボトルネックになる。

10000x10000配列,1000000x2回ランダムアクセス

リストオンリー
初期化：3.693434秒
ランダムアクセス：0.630347秒

numpy
初期化：0.000024秒
ランダムアクセス：0.979282秒

シングルディクショナリ
初期化：0.004809秒
ランダムアクセス：0.857782秒

ダブルディクショナリ
初期化：0.000006秒
ランダムアクセス：0.977852秒

↑ランダムアクセスが増えると、numpyの成績が悪くなってきた？

100000x100000配列,1000000x2回ランダムアクセス

numpy
初期化：0.000185秒
ランダムアクセス：6.354833秒

シングルディクショナリ
初期化：0.050961秒
ランダムアクセス1.063340秒

ダブルディクショナリ
初期化：0.000008秒
ランダムアクセス：1.238775秒

↑リストは初期化に時間がかかりすぎるので退場してもらった。
配列サイズが大きくなるとnumpyのランダムアクセスが極端に遅くなった。

結論（的なもの？）

• リストは初期化が遅いので、配列サイズが大きい場合は避けたほうが良い。

• numpyは要素へのアクセスが遅い傾向がありそう。また、初期化は速いがライブラリインポートに時間がかかるのでメリットは帳消し？（numpyライブラリインポート：0.072744秒）

• ディクショナリを使えば多分大丈夫そう。経路探索などでのランダムアクセスを考慮すると、やや速いシングルにしておいたほうがいいかも？（defaultdictライブラリインポート：0.002941秒）

• リストの初期化がもっともっと高速になってくれれば……

おまけ（PyPy）

pypy3.5-6.0.0 でもやってみました。

defaultdictライブラリインポート：0.000022秒
numpyライブラリインポート：0.217079秒
10000 1000000
10000x10000配列,1000000x2回ランダムアクセス

リストオンリー
初期化：1.210065秒
ランダムアクセス：0.081145秒

numpy
初期化：0.000100秒
ランダムアクセス：5.195924秒

シングルディクショナリ
初期化：0.012857秒
ランダムアクセス：1.842157秒

ダブルディクショナリ
初期化：0.000026秒
ランダムアクセス：1.217955秒

↑numpyのインポート時間が増えた。
PyPyではリストは速くなってるけど、numpy、ディクショナリは遅くなったかな？

defaultdictライブラリインポート：0.000030秒
numpyライブラリインポート：0.231340秒
100000 1000000
100000x100000配列,1000000x2回ランダムアクセス

numpy
初期化：0.000191秒
ランダムアクセス：12.990170秒

シングルディクショナリ
初期化：0.018306秒
ランダムアクセス：1.385379秒

ダブルディクショナリ
初期化：0.000021秒
ランダムアクセス：2.127735秒

計測に使用したコード

time.py

# -*- coding: utf-8 -*-
import time
import random
start = time.time()
from collections import defaultdict
elapsed_time = time.time() - start
print(f"defaultdictライブラリインポート：{elapsed_time:6f}秒")


start = time.time()
import numpy as np
elapsed_time = time.time() - start
print(f"numpyライブラリインポート：{elapsed_time:6f}秒")


N, M = map(int, input().split())    #N個の頂点のグラフ、M回ランダムアクセス
LRD = [[random.randrange(0, N), random.randrange(0, N),
        random.randrange(0, N)] for i in range(M)]

#N,M 説明
print(f"{N}x{N}配列,{M}x2回ランダムアクセス")
print()

# リストオンリー
if N <= 10000:
    print("リストオンリー")
    start = time.time()
    graph = [[0 for _ in range(N)] for _ in range(N)]
    elapsed_time = time.time() - start
    print(f"初期化：{elapsed_time:6f}秒")

    start = time.time()
    for L, R, D in LRD:
        graph[L][R] = D
        graph[R][L] = -D
    elapsed_time = time.time() - start
    print(f"ランダムアクセス：{elapsed_time:6f}秒")
    print()

    del start
    del graph

# numpy
print("numpy")
start = time.time()
graph = np.zeros((N,N),dtype=np.int)
elapsed_time = time.time() - start
print(f"初期化：{elapsed_time:6f}秒")

start = time.time()
for L, R, D in LRD:
    graph[L,R] = D
    graph[R,L] = -D
elapsed_time = time.time() - start
print(f"ランダムアクセス：{elapsed_time:6f}秒")
print()

del start
del graph


# シングルディクショナリ
print("シングルディクショナリ")
start = time.time()
graph = [defaultdict(int) for _ in range(N)]
elapsed_time = time.time() - start
print(f"初期化：{elapsed_time:6f}秒")

start = time.time()
for L, R, D in LRD:
    graph[L][R] = D
    graph[R][L] = -D
elapsed_time = time.time() - start
print(f"ランダムアクセス：{elapsed_time:6f}秒")
print()

del start
del graph

# ダブルディクショナリ
print("ダブルディクショナリ")
start = time.time()
graph = defaultdict(lambda: defaultdict(int))
elapsed_time = time.time() - start
print(f"初期化：{elapsed_time:6f}秒")

start = time.time()
for L, R, D in LRD:
    graph[L][R] = D
    graph[R][L] = -D
elapsed_time = time.time() - start
print(f"ランダムアクセス：{elapsed_time:6f}秒")
print()