はじめに

とあるプログラムを速くしてくれと頼まれて目につくものは大体どうにかしたのですがまだ遅い。いろいろ調べていると以下のようなコードに時間がかかっていることがわかりました。

def func(x, y, z):
    j = 0
    for i in range(len(x)):
        if y[i] == 1:
            x[i] = z[j]
            j = j + 1

気にはなっていたのですがx, yとzの要素数違うしNumPyでどうにかするの無理かなーと思ってそのままにしていました。しかし、

何気ない遅さがスピード狂のチューニング魂に火をつけた

funcの処理説明

funcが何をやっているのかをまず説明します。

xはどのような値でも構いません。0が並んだリストとします。

x = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

yとzは以下のような感じです。

y = [1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
z = [2, 3]

len(x) == len(y)という想定で、y[i]が1となるところのx[i]にzの値を前から詰めていきます。yの1の数とlen(z)は等しい必要があります。結果、func実行後のxは以下のようになります。

[2, 0, 0, 3, 0, 0, 0, 0, 0, 0]

今書いてても「こう書いた方が効率いいんじゃね？」というそもそものアルゴリズム変更が思いつきもしますが、ともかくこれをタイトルにあるようにNumPyのBoolean Indexを使って速くしたという話です。

改善後のコードと測定

こんな感じです。なお上ではx, y, zはlistとして説明しましたが実際にはNumPy配列です。

def func_improve(x, y, z):
    x[y == 1] = z

説明は後にしてとりあえず時間を計測します。あまり時間かからなそうなので¹1000回ループをさらに10回平均します。

import time
import numpy as np

x = [0] * 1000
y = [1] + [0] * 499 + [1] + [0] * 499
z = [1, 1]

x = np.array(x)
y = np.array(y)
z = np.array(z)

def perf(f):
    elapsed = []
    for _ in range(10):
        t1 = time.time()
        for _ in range(1000):
            f(x, y, z)
        t2 = time.time()
        elapsed.append(t2 - t1)
    print(f.__name__, np.average(elapsed))

perf(func)
perf(func_improve)

結果。50倍ぐらい速くなりました。

func 1.150579571723938
func_improve 0.023179507255554198

Boolean Indexと値設定

では何やっているかを説明します。ここで使っているのはBoolean Indexという技術です。
NumPyでは以下のようなことができます。

>>> a = np.arange(10)
>>> a[a % 2 == 0]
array([0, 2, 4, 6, 8])

a % 2 == 0を表示すると以下のようになります。つまり各要素について% 2 == 0を行い、結果の配列を作ります。

>>> a % 2 == 0
array([ True, False,  True, False,  True, False,  True, False,  True, False])

NumPyはIndexとしてbooleanの配列を受け取るとTrueのもののみ取り出した配列が返されます。

で、本題。あまり知られてない気がしますが（ってチュートリアル確認してたら書いてありますが）Boolean Indexは取り出しだけでなく設定も行えます。

>>> a[a % 2 == 0] = -1
>>> a
array([-1,  1, -1,  3, -1,  5, -1,  7, -1,  9])

つまり、条件式がTrueとなる要素について右辺の値を設定するわけです。

今回のコード改善のツボ

というわけで、Boolean Index使って値を設定することは知っていたのですが、今回は設定する値をスライドさせていかないといけないから無理かなーと思ってとりあえず書いてみたら、普通に動きました。

つまりこういうこと、

>>> a = np.arange(10)
>>> a[a % 2 == 0] = [-1, -2, -3, -4, -5]
>>> a
array([-1,  1, -2,  3, -3,  5, -4,  7, -5,  9])

条件式がTrueになる数と右辺の数は一致させる必要があります。まあ今回はそういう前提なので問題なし。一致しないとこんな感じにValueErrorになります。

>>> a = np.arange(10)
>>> a[a % 2 == 0] = [-1, -2, -3]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: NumPy boolean array indexing assignment cannot assign 3 input values to the 5 output values where the mask is true

まとめ

というわけで今回は配列の要素を設定するという処理をBoolean Indexを使って高速化しました。

Boolean Indexは取り出しだけでなく値の設定も行うことができます。さらに右辺は単一の値である必要はなく、Trueの数と同じになる必要はありますが配列を指定し順番に値を設定させるということもできます。
ここら辺を使うと要素設定を非常に高速に行わせることができます。ただし、はじめに書いたようにそもそものアルゴリズムを変えた方がいい気もしますが。

アルゴリズム変えてみた

「アルゴリズム変えた方がいいんじゃない」って書き逃げするのもあれなので、実際試してみました。

そもそも毎回yを回すのが遅いので要素が1となるインデックスを事前に拾っておきます。なおyを使いまわしてるのは上に書いてある計測用のperfを使いまわすためです。

y = [i for i in range(len(y)) if y[i]]

ｙがこういう配列だとfuncはこう書き直すことができる。ちなみにenumerate使ってyの参照少なくすると早くなるかなと思いましたがあまり変わらなかったです。

def func2(x, y, z):
    for i in range(len(y)):
        x[y[i]] = z[i]

NumPyのFancy Index使うとこう書けます。

def func2_improve(x, y, z):
    x[y] = z

で、計測結果。func2の方がfunc2_improveより速いのは予想外ですがやはりアルゴリズム自体を変えて計算量を減らすと大きな改善が見られることがわかりますね。

func 1.152672290802002
func_improve 0.020763278007507324
func2 0.004687786102294922
func2_improve 0.009325671195983886

1回の実行自体はそんなにかからないわけですが、コアの部分で何万回も呼び出されていたわけで、実行時間の大半を占めていたわけです。 ↩

Boolean Indexを使った要素設定の高速化