有点坑的 np.array_split

旧城等待， 2021-12-05 22:34 303阅读 0赞

假如我们想要把一个 array-like 的东西拆分成多个小组（或者叫 bucket），给定组的数目或者组的大小，最后一组不足则保持原样。

那么我们想到的可能是 `np.array_split()` 函数。我们使用其[官方例子][Link 1]来说明一下。

例子 1：

把一个含有 8 个元素的列表拆分成 3 组，那么 `np.array_split()` 的结果是：

>>> x = np.arange(8.0)
    >>> np.array_split(x, 3)
        [array([ 0.,  1.,  2.]), array([ 3.,  4.,  5.]), array([ 6.,  7.])]

例子 2：

上面的结果好像看起来没啥问题，那么如果列表中含有 7 个元素呢？

>>> x = np.arange(7.0)
    >>> np.array_split(x, 3)
        [array([ 0.,  1.,  2.]), array([ 3.,  4.]), array([ 5.,  6.])]

结果每组的大小是 322，很明显不是我们想要的结果，我们想要的结果是 331。

再次查看[官方文档][Link 1]，可以看到如下关于该函数如何拆分的解释：

> 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.

也就是说，对于一个长度为 `l` 的数组，该函数是这么拆的：**前 `l % n` 个组的大小是 `l // n + 1`，剩下组的大小是 `l // n`。** 其中 `//` 表示下取整，即 `np.floor()`。

对于例子 1，`l` 为 8，`n` 为 3，前 `8 % 3 = 2` 个组的大小为 `8 // 3 + 1 = 3`，剩下组的大小为 `8 // 3 = 2`。

对于例子 2，`l` 为 7，`n` 为 3，前 `7 % 3 = 1` 个组的大小为 `7 // 3 + 1 = 3`，剩下组的大小为 `7 // 3 = 2`。

所以，从直观上来讲，该函数倾向于使得拆分后的组大小分布更为「均衡」，有时候可能不是我们想要的结果，值得注意。

另外，还有一个和 `np.array_split` 比较相似的函数：[`np.split()`][np.split]，这两个函数的唯一区别是 `split()` 必须是**等分**，否则会抛出异常：`ValueError: array split does not result in an equal division`。

## Reference ##

*  [numpy.array\_split — NumPy v1.16 Manual][Link 1]
 *  [numpy.split — NumPy v1.16 Manual][np.split]

## END ##

[Link 1]: https://docs.scipy.org/doc/numpy/reference/generated/numpy.array_split.html
[np.split]: https://docs.scipy.org/doc/numpy/reference/generated/numpy.split.html