numpyで全てのペアを取得する方法

n個の数字の全てのペアが欲しいとき、Pythonではitertoolsを使う方法が有名ですが、numpyのtriu_indices関数でも可能です。itertoolsをimportしなくて済むという利点があります。

例

コード

import numpy as np
print(np.array(np.triu_indices(5,1)).T)  # 第2引数は1にする

出力結果

[[0 1]
 [0 2]
 [0 3]
 [0 4]
 [1 2]
 [1 3]
 [1 4]
 [2 3]
 [2 4]
 [3 4]]

参考：一般的なitertoolsを使う方法

コード

import itertools
print(list(itertools.combinations(range(5),2)))

出力結果

[(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]

説明

triu_indicesは行列の上三角部分の要素の行番号と列番号を返す関数です。第1引数が行列の大きさ、第2引数は0だと対角項を含み、1だと対角項を含みません。ここでは異なる2要素の組み合わせが欲しいので、第2引数は1にして下さい。

triu_indicesの元の出力は以下のような感じです。このままだと使いづらいので、上記の例のようにnp.arrayに変換して転置すると良いです。

print(np.triu_indices(5,1))

出力結果

(array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([1, 2, 3, 4, 2, 3, 4, 3, 4, 4]))

補足

ただの数字でなく、何らかのリストの要素のペアが欲しい場合は、以下のように対象リストをnp.array化して"[ ]"の中に要素番号を入れれば良いです。

x = ['A', 'B', 'C', 'D', 'E']
c = np.array(np.triu_indices(len(x),1)).T
print(np.array(x)[c])

出力結果

[['A' 'B']
 ['A' 'C']
 ['A' 'D']
 ['A' 'E']
 ['B' 'C']
 ['B' 'D']
 ['B' 'E']
 ['C' 'D']
 ['C' 'E']
 ['D' 'E']]