避免for循环使用numpy通过索引值将数组拆分为多个数组答案

【问题标题】：Avoid for-loop to split array into multiple arrays by index values using numpy避免for循环使用numpy通过索引值将数组拆分为多个数组
【发布时间】：2021-03-22 02:49:35
【问题描述】：

输入：有两个输入数组：

value_array = [56, 10, 65, 37, 29, 14, 97, 46]
index_array = [ 0,  0,  1,  0,  3,  0,  1,  1]

输出：我想在不使用 for 循环的情况下使用 index_array 拆分 value_array。所以输出数组将是：

split_array = [[56, 10, 37, 14],  # index 0
               [65, 97, 46],      # index 1
               [],                # index 2
               [29]]              # index 3

有什么方法可以在不使用任何 for 循环的情况下使用 numpy 来做到这一点？我看过numpy.where，但不知道该怎么做。

For 循环：这是使用 for 循环的方法。我想避免for循环。

split_array = []
for i in range(max(index_array) + 1):
    split_array.append([])
           
for i in range(len(value_array)):
    split_array[index_array[i]].append(value_array[i])

【问题讨论】：

你有一个带有循环的工作版本吗？
@ivan 添加了带有 for 循环的代码。
由于 split_array 是一个包含不同长度的数组/列表的列表，因此您无法避免 Python 级别的循环（即使它隐藏在 np.split 函数中）。传说中的no-loop 操作仅限于整个数组（或切片）并返回整个（数字）数组，而不是列表列表。

标签： python arrays numpy for-loop

【解决方案1】：

这就够了吗？

方案一（注意：for循环不是遍历整个索引数组）

import numpy as np

value_array = np.array([56, 10, 65, 37, 29, 14, 97, 46])
index_array = np.array([ 0,  0,  1,  0,  3,  0,  1,  1])

max_idx = np.max(index_array)
split_array = []

for idx in range(max_idx + 1):
    split_array.append([])
    split_array[-1].extend(list(value_array[np.where(index_array == idx)]))
print(split_array)

[[56, 10, 37, 14], [65, 97, 46], [], [29]]

解决方案 2

import numpy as np

value_array = np.array([56, 10, 65, 37, 29, 14, 97, 46])
index_array = np.array([ 0,  0,  1,  0,  3,  0,  1,  1])

value_array = value_array[index_array.argsort()]
split_idxs = np.squeeze(np.argwhere(np.diff(np.sort(index_array)) != 0) + 1)
print(np.array_split(value_array, split_idxs))

[array([56, 10, 37, 14]), array([65, 97, 46]), array([29])]

【讨论】：

更新了我的问题。我只是想通过利用numpy 来避免循环。编辑：我注意到，您的代码有 max_idx 迭代的 for 循环，而不是 len(value_array)，这很好。
解决方案二的结果缺少一个空数组，对吧？

【解决方案2】：

你可以这样做：

import numpy as np

value_array = np.array([56, 10, 65, 37, 29, 14, 97, 46])
index_array = np.array([ 0,  0,  1,  0,  3,  0,  1,  1])

# find the unique values in index array and the corresponding counts
unique, counts = np.unique(index_array, return_counts=True)

# create an array with 0 for the missing indices
zeros = np.zeros(index_array.max() + 1, dtype=np.int32)
zeros[unique] = counts  # zeros = [4 3 0 1] 0 -> 4, 1 -> 3, 2 -> 0, 3 -> 1

# group by index array
so = value_array[np.argsort(index_array)]  # so = [56 10 37 14 65 97 46 29]

# finally split using the counts 
res = np.split(so, zeros.cumsum()[:-1])

print(res)

输出

[array([56, 10, 37, 14]), array([65, 97, 46]), array([], dtype=int64), array([29])]

这种方法的时间复杂度是O(N logN)。

此外，如果您不关心缺少的索引，您可以使用以下内容：

_, counts = np.unique(index_array, return_counts=True)
res = np.split(value_array[np.argsort(index_array)], counts.cumsum()[:-1])

print(res)

输出

[array([56, 10, 37, 14]), array([65, 97, 46]), array([29])]

【讨论】：

np.split 仍然循环 - 列表中每个返回的数组一个切片。
@hpaulj 很高兴知道！

【解决方案3】：

确实，您可以通过数组来使用 numpy：

import numpy as np
value_array=np.array(value_array)
index_array=np.array(index_array)
split_array=[value_array[np.where(index_array==j)[0]] for j in set(index_array)]

【讨论】：