如何使用 python3 将“for”循环转换为列表列表的矩阵表达式？

Question

我需要将for 循环转换为使用矩阵形式的表达式。我有一个列表列表、一个索引列表和一个名为“toSave”的形状 (4,2) 矩阵：

import numpy as np

M = [list() for i in range(3)]
indices= [1,1,0,1]
toSave = np.array([[0, 0],
                   [0, 1],
                   [0, 2],
                   [0, 3]])

对于索引中的每个索引i，我想保存与索引中索引i 的位置相对应的行：

for n, i in enumerate(indices):
    M[i].append(toSave[n])

结果是：

M=[[[0, 2]], [[0, 0], [0, 1], [0, 3]], []]

是否可以使用矩阵表达式来代替 for 循环，例如 M[indices].append(toSave[range(4)]) ？

Answer 1

这是一种方法 -

sidx = np.argsort(indices)
s_indx = np.take(indices, sidx)

split_idx = np.flatnonzero(s_indx[1:] != s_indx[:-1])+1
out = np.split(toSave[sidx], split_idx, axis=0)

示例运行 -

# Given inputs
In [67]: M=[[] for i in range(3)]
    ...: indices= [1,1,0,1]
    ...: toSave=np.array([[0, 0],
    ...:        [0, 1],
    ...:        [0, 2],
    ...:        [0, 3]])
    ...: 

# Using loopy solution
In [68]: for n, i in enumerate(indices):
    ...:     M[i].append(toSave[n])
    ...:     

In [69]: M
Out[69]: [[array([0, 2])], [array([0, 0]), array([0, 1]), array([0, 3])], []]

# Using proposed solution
In [70]: out
Out[70]: 
[array([[0, 2]]), array([[0, 0],
        [0, 1],
        [0, 3]])]

性能提升

一种更快的方法是避免np.split 并使用slicing 进行拆分，就像这样 -

sorted_toSave = toSave[sidx]
idx = np.concatenate(( [0], split_idx, [toSave.shape[0]] ))
out = [sorted_toSave[i:j] for i,j in zip(idx[:-1],idx[1:])]