Selecting items on a matrix based on indexes given by an array-CodePudding

Consider this matrix:

[0.9, 0.45, 0.4, 0.35],
[0.4, 0.8, 0.3, 0.25],
[0.5, 0.45, 0.9, 0.35],
[0.2, 0.18, 0.8, 0.1],
[0.6, 0.45, 0.4, 0.9]

and this list:

[0,1,2,3,3]

I want to create a list that looks like the following:

[0.9, 0.8, 0.9, 0.1, 0.9]

To clarify, for each row, I want the element of the matrix whose column index is contained in the first array. How can I accomplish this?

CodePudding user response：

Zip the two lists together as below

a=[[0.9, 0.45, 0.4, 0.35],[0.4, 0.8, 0.3, 0.25],[0.5, 0.45, 0.9, 0.35],[0.2, 0.18, 0.8, 0.1],[0.6, 0.45, 0.4, 0.9]]
b=[0,1,2,3,3]
[i[j] for i,j in zip(a,b)]

Result

[0.9, 0.8, 0.9, 0.1, 0.9]

This basically pairs up each sublist in the matrix with the element of your second list in order with zip(a,b)

Then for each pair you choose the bth element of a

CodePudding user response：

If this is a numpy array, you can pass in two numpy arrays to access the desired indices:

import numpy as np

data = np.array([[0.9, 0.45, 0.4, 0.35],
[0.4, 0.8, 0.3, 0.25],
[0.5, 0.45, 0.9, 0.35],
[0.2, 0.18, 0.8, 0.1],
[0.6, 0.45, 0.4, 0.9]])

indices = np.array([0,1,2,3,3])

data[np.arange(data.shape[0]), indices]

This outputs:

[0.9 0.8 0.9 0.1 0.9]

CodePudding user response：

In the first array [0, 1, 2, 3, 3], the row is determined by the index of the each element, and the value at that index is the column. This is a good case for enumerate:

matrix = [[ ... ], [ ... ], ...] # your matrix
selections = [0, 1, 2, 3, 3]
result = [matrix[i][j] for i, j in enumerate(selections)]

This will be much more efficient than looping through the entire matrix.

CodePudding user response：

Loop through both arrays together using the zip function.

def create_array_from_matrix(matrix, indices):
    if len(matrix) != len(indices):
        return None
    res = []
    for row, index in zip(matrix, indices):
        res.append(row[index])

    return res