静かなる名辞

pythonとプログラミングのこと



【python】numpy配列を分割する

はじめに

 numpy配列を分割したくなることがたまにある。

 当然というかさすがというか、それ用の関数が用意されている。でも使い方をよく忘れるので覚書として書いておく。

 目次

np.split

 そのまんますぎる名前の関数がある。

numpy.split — NumPy v1.15 Manual

>>> import numpy as np
>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.split(a, 2)
[array([0, 1, 2, 3, 4]), array([5, 6, 7, 8, 9])]

 わかりやすい。返り値は「リスト」で、ドキュメントにも明確にそう書いてある。

 第二引数は「indices_or_sections : int or 1-D array」である。配列も渡せるので試してみる。

>>> np.split(a, [2,8])
[array([0, 1]), array([2, 3, 4, 5, 6, 7]), array([8, 9])]

 要するにスライスですね。こうしたのと一緒。

>>> [a[:2], a[2:8], a[8:]]
[array([0, 1]), array([2, 3, 4, 5, 6, 7]), array([8, 9])]

 numpyらしくaxisも指定できる。デフォルトは0。これは2次元配列で試してみよう。

>>> b = np.arange(16).reshape(4, 4)
>>> b
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])
>>> np.split(b, 2)
[array([[0, 1, 2, 3],
       [4, 5, 6, 7]]),
 array([[ 8,  9, 10, 11],
       [12, 13, 14, 15]])]
>>> np.split(b, 2, axis=1)
[array([[ 0,  1],
       [ 4,  5],
       [ 8,  9],
       [12, 13]]),
 array([[ 2,  3],
       [ 6,  7],
       [10, 11],
       [14, 15]])]
>>> np.split(b, [1,3])
[array([[0, 1, 2, 3]]),
 array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]]),
 array([[12, 13, 14, 15]])]
>>> np.split(b, [1,3], axis=1)
[array([[ 0],
       [ 4],
       [ 8],
       [12]]),
 array([[ 1,  2],
       [ 5,  6],
       [ 9, 10],
       [13, 14]]),
 array([[ 3],
       [ 7],
       [11],
       [15]])]

 結果が見づらいので、適当に整形してある。見ての通りの使い方ができる。

np.array_split

 array_splitというちょっと名前が違う関数がある。

numpy.array_split — NumPy v1.15 Manual

 ドキュメントを見ると、存在意義を疑いたくなるようなことが書いてある。

Please refer to the split documentation. The only difference between these functions is that array_split allows indices_or_sections to be an integer that does not equally divide the axis. For an array of length l that should be split into n sections, it returns l % n sub-arrays of size l//n + 1 and the rest of size l//n.

>>> a = np.arange(10)
>>> np.split(a, 3)
Traceback (most recent call last):
  File "***/lib/python3.5/site-packages/numpy/lib/shape_base.py", line 553, in split
    len(indices_or_sections)
TypeError: object of type 'int' has no len()

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "***/lib/python3.5/site-packages/numpy/lib/shape_base.py", line 559, in split
    'array split does not result in an equal division')
ValueError: array split does not result in an equal division
>>> np.array_split(a, 3)
[array([0, 1, 2, 3]), array([4, 5, 6]), array([7, 8, 9])]

 たったこれだけの違い。もう少し細かく見ていくと、

>>> np.array_split(np.arange(11), 5)
[array([0, 1, 2]), array([3, 4]), array([5, 6]), array([7, 8]), array([ 9, 10])]
>>> np.array_split(np.arange(12), 5)
[array([0, 1, 2]), array([3, 4, 5]), array([6, 7]), array([8, 9]), array([10, 11])]
>>> np.array_split(np.arange(13), 5)
[array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]), array([ 9, 10]), array([11, 12])]
>>> np.array_split(np.arange(14), 5)
[array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]), array([ 9, 10, 11]), array([12, 13])]

 前の方の配列から長さが増えていく感じの挙動になる。

 ちなみに、他の使い方はすべて同じで、第二引数に配列を渡したりaxisを変えたりもできる。

 同じ関数にしてオプション引数で挙動を変えれば十分な気がするが、歴史的経緯などで統一されていないのかなぁ。

まとめ

 特にまとめることもありませんが、そんな感じでできます。

 割り切れる場合はnp.split()、割り切れない場合はnp.array_split()を使うのかな? とも思ったが、どちらでもいけるnp.array_split()だけ覚えておけば良いような気もする。お好みで。