今日覚えたNumPyのテクニック

今日覚えたNumPyのテクニックを紹介します。

numpy.stack で2次元配列を繋げて3次元配列を作る

1次元配列を繋げて2次元配列を作りたいような場合には numpy.vstack や numpy.hstack という関数があります。

a = numpy.array([1, 2, 3, 4])
b = numpy.array([2, 4, 6, 8])
c = numpy.array([3, 6, 9, 12])
numpy.vstack([a, b, c])

# array([[ 1,  2,  3,  4],
#        [ 2,  4,  6,  8],
#        [ 3,  6,  9, 12]])

2次元配列を繋げて3次元配列を作りたいような場合にはどのようにすればよいか知らなかったのですが，numpy.stack という関数があることを知りました。

a = numpy.array([[1, 2, 3], [4, 5, 6]])
b = numpy.array([[2, 4, 6], [8, 10, 12]])
c = numpy.array([[3, 6, 9], [12, 15, 18]])
numpy.stack([a, b, c])

# array([[[ 1,  2,  3],
#         [ 4,  5,  6]],
#        [[ 2,  4,  6],
#         [ 8, 10, 12]],
#        [[ 3,  6,  9],
#         [12, 15, 18]]])

条件が3つ以上ある場合の条件演算子を integer array indexing で実現する

条件演算子 x ? a : b のようなことを配列のすべての要素に対してやりたい場合には numpy.where を使うことができます。

x = numpy.array([0, 1, 0, 1])
a = numpy.array([1, 2, 3, 4])
b = numpy.array([5, 6, 7, 8])
numpy.where(x, a, b)

# array([5, 2, 7, 4])

条件が3つ以上ある場合にどうすればよいか知らなかったので，いろいろ考えました。最初に思いつくのは numpy.where をネストして無理やり実現する方法です。

x = numpy.array([0, 1, 2, 1])
a = numpy.array([1, 2, 3, 4])
b = numpy.array([5, 6, 7, 8])
c = numpy.array([9, 10, 11, 12])
numpy.where(x == 0, a, numpy.where(x == 1, b, c))

# array([ 1,  6, 11,  8])

この方法では場合分けの個数が変化する場合に対応できなくなってしまうので不便です。別の方法として NumPy の indexing を使う方法を教えてもらいました。

x = numpy.array([0, 1, 2, 1])
a = numpy.array([1, 2, 3, 4])
b = numpy.array([5, 6, 7, 8])
c = numpy.array([9, 10, 11, 12])
u = numpy.stack([a, b, c])

# array([[ 1,  2,  3,  4],
#        [ 5,  6,  7,  8],
#        [ 9, 10, 11, 12]])

u[x, numpy.arange(x.shape[0])]

# array([ 1,  6, 11,  8])