1次元のnumpy.ndarrayから対となる組み合わせを生成する効率的な方法

Fast (numpy) pairwise combination generation from 1-dimensional array (pandas index) – StackOverflow

組み合わせは直積集合から上三角行列を求めるイメージ.

import numpy as np


def numpy_combinations(arr):
    r, c = np.triu_indices(len(arr), 1)
    return np.stack((arr[r], arr[c]), 1)


a = np.array([1, 2, 3, 4])
numpy_combinations(a)
array([[1, 2],
       [1, 3],
       [1, 4],
       [2, 3],
       [2, 4],
       [3, 4]])
from numba import jit


@jit(nopython=True, parallel=True, fastmath=True)
def pairwise_combs(a):
    n = len(a)
    L = n*(n-1)//2
    out = np.empty((L,2),dtype=a.dtype)
    iterID = 0
    for i in range(n):
        for j in range(i+1,n):
            out[iterID,0] = a[i]
            out[iterID,1] = a[j]
            iterID += 1
    return out
a = np.random.rand(1000)
%timeit pairwise_combs(a)
%timeit numpy_combinations(a)
a = np.random.rand(5000)
%timeit pairwise_combs(a)
%timeit numpy_combinations(a)
100 loops, best of 3: 3.63 ms per loop
100 loops, best of 3: 18.1 ms per loop
10 loops, best of 3: 47.8 ms per loop
1 loop, best of 3: 336 ms per loop
カテゴリー: 未分類 パーマリンク

コメントを残す

以下に詳細を記入するか、アイコンをクリックしてログインしてください。

WordPress.com ロゴ

WordPress.com アカウントを使ってコメントしています。 ログアウト /  変更 )

Google フォト

Google アカウントを使ってコメントしています。 ログアウト /  変更 )

Twitter 画像

Twitter アカウントを使ってコメントしています。 ログアウト /  変更 )

Facebook の写真

Facebook アカウントを使ってコメントしています。 ログアウト /  変更 )

%s と連携中

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください