有没有更快的方法从多维数组中检索数组，使用索引数组？

Question

我正在尝试加速一段被称为LOT的代码，以期减少脚本运行时间。

假设我有一个多维数组：

[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

和一维指数数组：

[2], [0], [1]

没有循环，有没有办法从多维数组中检索这些索引，即：

[3], [4], [8]

任何帮助赞赏！

Answer 1

你可以使用itertools.starmap

import itertools


def get_values_from_indices(array_values, array_indices):
    """
    This function will accept two params, 
    once is a multi-dimensional list, and other one is list of indices.
    """
    return list(itertools.starmap(lambda x, y: x[y[0]], zip(array_values, array_indices)))

DEMO

multi_dimensional_array = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
list_of_indices = [[2], [0], [1]]

result = get_values_from_indices(multi_dimensional_array , list_of_indices)

print(result)
# [3, 4, 8]

Answer 2

不确定你将如何想要你的结果，但有了numpy你可以实现类似下面的东西。

import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# array([[1, 2, 3],
#   [4, 5, 6],
#   [7, 8, 9]])

a[[0,1,2],[2,0,1]] # array([3, 4, 8])
a[[0,1,2],[1,2,1]] # array([2, 6, 8])

甚至可能是，

indices = [[2],[0],[1]]
a[range(len(indices)), np.reshape(indices, -1)] # array([3, 4, 8])

Answer 3

列表理解：

listOfList = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

indexInto = [2, 0, 1]       # first version works for this
indexTwo = [[2], [0], [1]]  # second version works for this

# first version
values = [listOfList[lili][indexInto[lili]] for lili in range(len(listOfList))] # both lists need same length

# second version
values2 = [listOfList[lili][indexTwo[lili][0]] for lili in range(len(listOfList))] # both lists need same length

print( values)
print( values2)

输出：

[3, 4, 8]    
[3, 4, 8]

Answer 4

另一个numpy解决方案：

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

a2 = [[2], [0], [1]]

a[np.arange(len(a)), np.concatenate(a2)]  # array([3, 4, 8])

Answer 5

Numpy解决方案：

import numpy as np

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
ind = [[2], [0], [1]]
a = np.array(L)
b = np.array(ind)

c = a[np.arange(len(a)), b.reshape(-1)]
print (c.tolist())
[3, 4, 8]

Answer 6

带有映射的Lambda解决方案，无需在一行中导入任何外部模块：

list_1=[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
indices=[[2], [0], [1]]

print(list(map(lambda x,y :list(map(lambda z:x[z],y)),list_1,indices)))

输出：

[[3], [4], [8]]