时间戳

首页 > 解决方案 > 避免for循环使用numpy通过索引值将数组拆分为多个数组

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

问题描述

输入：有两个输入数组：

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

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

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

有没有办法在numpy不使用任何for循环的情况下做到这一点？我已经看过numpy.where但无法弄清楚如何做到这一点。

For-loop：这是使用 for-loop 的方法。我想避免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])

标签： pythonarraysnumpyfor-loop

解决方案

你可以这样做：

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])]

推荐阅读