基于另一个列表的组列表元素

时间:2016-09-08 09:35:00

标签: python arrays numpy

我有两个列表:inpbase

我想根据inp中的位置,将out中的每个项目添加到base的列表中。

以下代码可以正常使用:

from pprint import pprint as print

num = 3
out = [[] for i in range(num)]
inp = [[1,1],[2,1],[3,2],[7,11],[9,99],[0,-1]]
base = [0,1,0,2,0,1]

for i, num in enumerate(base):
    out[num].append(inp[i])
print(out,width=40)

[[[1, 1], [3, 2], [9, 99]],
 [[2, 1], [0, -1]],
 [[7, 11]]]

我想使用NumPy模块(np.arraynp.append等)来完成此操作。

任何人都可以帮助我吗?

1 个答案:

答案 0 :(得分:1)

假设baseinp为NumPy数组,我们可以这样做 -

# Get sorted indices for base
sidx = base.argsort()

# Get where the sorted version of base changes groups
split_idx = np.flatnonzero(np.diff(base[sidx])>0)+1
# OR np.unique(base[sidx],return_index=True)[1][1:]

# Finally sort inp based on the sorted indices and split based on split_idx
out = np.split(inp[sidx], split_idx)

为了使它适用于列表,我们需要进行一些调整,主要是索引部分,我们可以使用np.take将索引替换为前面方法中列出的数组。因此,修改后的版本将是 -

sidx = np.argsort(base)
split_idx = np.flatnonzero(np.diff(np.take(base,sidx))>0)+1
out = np.split(np.take(inp,sidx,axis=0), split_idx)