I'm trying to speed up a section of code that's called a LOT in the hope to cut a script run-time down.
Say I have a multidimensional array:
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
And a single-dimensional array of indices:
[2], [0], [1]
Without a loop, is there a way to retrieve those indices from the multi-dimensional array, i.e,:
[3], [4], [8]
Any help appreciated!
You could use 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]
Not sure exactly how you would want your result, but with numpy you could achieve something similar below.
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])
May even be,
indices = [[2],[0],[1]]
a[range(len(indices)), np.reshape(indices, -1)] # array([3, 4, 8])
List comprehensions:
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)
Output:
[3, 4, 8]
[3, 4, 8]
Another numpy solution:
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])
Numpy solution:
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]
Lambda solution with map without importing any external module in just one line :
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)))
output:
[[3], [4], [8]]
