zip()を使って複数配列(array, list)をshuffleする(python)

複数のリストやnumpy arrayを、その対応関係を崩さずにシャッフルするのに手間取ったのでメモする。ここでいう対応関係とは、配列Xとyがあった時、(X[0], y[0]), (X[1], y[1])のペアのことを表す。一つの訓練サンプルに対するラベルとか、従属変数に対する説明変数とか、そんな感じ。

今まではarray生成してindexの乱数を割り振って入れ直していた。 下みたいな感じ

# 配列X, yを与え、対応を崩さずにシャッフルする
def shuffle_samples(X, y):
    order = np.arange(X.shape[0])
    np.random.shuffle(order)
    X_result = np.zeros(X.shape)
    y_result = np.zeros(y.shape)
    for i in range(X.shape[0]):
        X_result[i, ...] = X[order[i], ...]
        y_result[i, ...] = y[order[i], ...]
    return X_result, y_result

np.arange()をシャッフルして、その順番で入力配列を入れ替える。出力はnp.array二つ。なんとなく回りくどい気がする。

そんな時zip()を使うと複数配列を同時に扱いやすいことを知った

a = [1, 2, 3, 4, 5, 6]
b = ['a', 'b', 'c', 'd', 'e', 'f']

zip(a, b)
# <zip at 0x10f721188>
# zipオブジェクトが返される(中身は見えない)

l = list(zip(a, b))
# [(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd'), (5, 'e'), (6, 'f')]
# リストに変換すると見える。タプルのリストになっている

np.random.shuffle(l)     # シャッフル
# [(4, 'd'), (5, 'e'), (2, 'b'), (1, 'a'), (3, 'c'), (6, 'f')]

new_a, new_b = zip(*l)
# new_a  ->  (4, 5, 2, 1, 3, 6)
# new_b  ->  ('d', 'e', 'b', 'a', 'c', 'f')
# 型がタプルになっていることに注意

zip()で配列をまとめ、zip(*)で配列を展開(?)できる。shuffleの部分をソートにすれば複数配列ソートも実現できる。

これを使って先ほどの関数を書き換えると、

def shuffle_samples(X, y):
    zippded = list(zip(X, y))
    np.random.shuffle(zipped)
    X_result, y_result = zip(*zipped)
    return np.asarray(X_result), np.asarray(y_result)    # 型をnp.arrayに変換

すごく簡潔に書けた。zip(*)で展開すると配列の型がタプルになるので、出力前に変換している。

さらにzip()は入力が3つ以上でも大丈夫なので、任意の数の配列を受け取るように関数を書き換えると

def shuffle_samples(*args):
    # *argsで可変長引数を受け取る。変数argsにリストで格納される

    # unzipで複数配列のリスト -> 要素毎にまとめたタプルのリスト　に変換
    zipped = list(zip(*args))
    np.random.shuffle(zipped)

    # unzipして複数配列のリストの形に戻す
    shuffled = list(zip(*zipped))
    
    result = []
    # np.arrayに変換する処理
    for ar in shuffled:
        result.append(np.asarray(ar))
    return result

となる。

zip便利、というかpython便利。