如何按条件值指定的子数组长度分割分类numpy数组？

我需要按条件值分割NumPy数组，要求每个子数组的长度等于该条件值。例如：

import numpy as np

x = np.array(['Horses', 'Cats', 'Dogs'])
conditional_value = 1

期望输出：

array([['Horses'],
       ['Cats'],
       ['Dogs']], dtype='<U6')

我尝试用np.split方法，但结果不符合预期：

import numpy as np

x = np.array(['Horses', 'Cats', 'Dogs'])
np.split(x, 1)

实际输出：[array(['Horses', 'Cats', 'Dogs'], dtype='<U6')]

请问还有什么方法可以实现按条件值指定的子数组长度分割上述NumPy数组？

解决方案 #

方法1：使用np.reshape（通用方案） #

np.reshape是最直接的方法，只要原数组的总长度能被conditional_value整除，就能直接调整成目标形状。其中-1表示让NumPy自动计算该维度的元素个数，保证总元素数和原数组一致：

import numpy as np

x = np.array(['Horses', 'Cats', 'Dogs'])
conditional_value = 1
result = x.reshape(-1, conditional_value)
print(result)

输出：

array([['Horses'],
       ['Cats'],
       ['Dogs']], dtype='<U6')

如果条件值不是1，比如原数组长度为6、条件值为2，同样适用：

x = np.array(['a','b','c','d','e','f'])
conditional_value = 2
result = x.reshape(-1, conditional_value)
print(result)

输出：

array([['a', 'b'],
       ['c', 'd'],
       ['e', 'f']], dtype='<U1')

方法2：使用np.expand_dims（仅适用于子数组长度为1的场景） #

如果你的conditional_value固定为1，只是想给原一维数组增加一个维度，也可以用这个方法：

import numpy as np

x = np.array(['Horses', 'Cats', 'Dogs'])
result = np.expand_dims(x, axis=1)
print(result)

输出和方法1一致。

为什么np.split不适用？ #

np.split的第二个参数是分割点数量或分割索引列表，并非子数组长度。比如np.split(x, 1)是将数组分成1份（即原数组本身）；如果要按长度1分割，需要传入分割索引[1,2]，但得到的是多个一维数组的列表，而非二维数组：

np.split(x, [1,2])
# 输出：[array(['Horses'], dtype='<U6'), array(['Cats'], dtype='<U6'), array(['Dogs'], dtype='<U6')]

这和你需要的二维数组结构不符，所以不适合你的需求。

内容的提问来源于stack exchange，提问作者user3046211