おいも貴婦人ブログ

生物系博士課程満期退学をしたAIエンジニアのブログ。

numpy.arrayによる2次元配列から任意の列を抽出する方法

2次元配列を扱っていると、ある列だけ取り出したいことはよくあることです。そんなときに、np.arrayを使いましょう。
対象と成る配列は以下のようになります。

mat =\
[[ 0  1  2  3  4  5  6  7  8  9]
 [10 11 12 13 14 15 16 17 18 19]
 [20 21 22 23 24 25 26 27 28 29]
 [30 31 32 33 34 35 36 37 38 39]
 [40 41 42 43 44 45 46 47 48 49]
 [50 51 52 53 54 55 56 57 58 59]
 [60 61 62 63 64 65 66 67 68 69]
 [70 71 72 73 74 75 76 77 78 79]
 [80 81 82 83 84 85 86 87 88 89]
 [90 91 92 93 94 95 96 97 98 99]]
まずは2列目だけを取り出したい場合。
mat[:,2]

出力

array([ 2, 12, 22, 32, 42, 52, 62, 72, 82, 92])
2列目を2次元配列で取り出す場合。
mat[:,2:3]

出力

array([[ 2],
       [12],
       [22],
       [32],
       [42],
       [52],
       [62],
       [72],
       [82],
       [92]])
奇数列のみを取り出したい場合。
mat[:,1::2]

出力

array([[ 1,  3,  5,  7,  9],
       [11, 13, 15, 17, 19],
       [21, 23, 25, 27, 29],
       [31, 33, 35, 37, 39],
       [41, 43, 45, 47, 49],
       [51, 53, 55, 57, 59],
       [61, 63, 65, 67, 69],
       [71, 73, 75, 77, 79],
       [81, 83, 85, 87, 89],
       [91, 93, 95, 97, 99]])

正直、ここまではpythonを知っている人ならば当たり前過ぎる話です。

任意の列を取り出したい場合。(1列目、7列目、8列目、8列目)
mat[:,np.array([0,6,7,7])]

出力

array([[ 0,  6,  7,  7],
       [10, 16, 17, 17],
       [20, 26, 27, 27],
       [30, 36, 37, 37],
       [40, 46, 47, 47],
       [50, 56, 57, 57],
       [60, 66, 67, 67],
       [70, 76, 77, 77],
       [80, 86, 87, 87],
       [90, 96, 97, 97]])

Indexの代わりにarrayを代入できるらへんが面白いですね。