使用索引数组将numpy数组分成多个数组（Python）

Question

我有一个数组：

a = [1, 3, 5, 7, 29 ... 5030, 6000]

这个数组是由之前的一个过程创建的，数组的长度可能不同（这取决于用户的输入）。

我还有一个数组：

b = [3, 15, 67, 78, 138]

（这也可能完全不同）

我想使用数组b将数组a切成多个数组。

更具体地说，我希望结果数组为：

array1 = a[:3]
array2 = a[3:15]
...
arrayn = a[138:]

当 n = len(b) 时。

我的第一个想法是创建一个二维数组 slices，其维度为 (len(b), something)。然而，我们事先不知道这个 something 的值，因此我将其分配为 len(a)，因为那是它可能包含的最大数字数量。

我有这段代码：

 slices = np.zeros((len(b), len(a)))

 for i in range(1, len(b)):
     slices[i] = a[b[i-1]:b[i]]

但是我遇到了这个错误：

ValueError: could not broadcast input array from shape (518) into shape (2253412)

- pavlos163

---

我很惊讶我找不到重复的问题。干得好！ - Mad Physicist

Answer 1

您可以使用 numpy.split 函数来进行分割：

np.split(a, b)

示例：

np.split(np.arange(10), [3,5])
# [array([0, 1, 2]), array([3, 4]), array([5, 6, 7, 8, 9])]

- Psidom

Answer 2

b.insert(0,0)
result = []
for i in range(1,len(b)):
    sub_list = a[b[i-1]:b[i]]
    result.append(sub_list)
result.append(a[b[-1]:])

- Bohdan

---

这基本上就是 np.split 所做的事情，只不过稍微花了一些心思来允许在其他轴上进行分割。 - hpaulj

Answer 3

你之所以会出现这个错误，是因为你试图创建一个不规则的数组。在numpy中，这是不允许的。

@Bohdan的回答有所改进：

from itertools import zip_longest
result = [a[start:end] for start, end in zip_longest(np.r_[0, b], b)]

这里的技巧在于zip_longest使得最终的切片从b[-1]到None，相当于a[b[-1]:]，无需特殊处理最后一个元素。

请不要选择这个。这只是我为了好玩添加的东西。 “正确”的答案是@Psidom的答案。

- Mad Physicist