CODEWITHSUNDEEP
pythonarrayspython-2.7arraylistindexing
Mechanician
Mechanician

Reputation: 545

How do I efficiently find the row indices of multiple elements in a 2D matrix?

say,

C = [ [1,2,3],[1,3,4],[3,5,6]]
item_list=[1,3,4]

I used the following code to accomplish what I wanted :

    rind = [[i for i in range(len(C)) if item in C[i]]
        for item in item_list]

I got the rind to be [[0, 1], [0, 1, 2], [1]]

I actually want my o/p to be as a 1d array like [0 1 0 1 2 1]

Could you either suggest a completely alternative approach to obtain row indices or advise me on how to convert the list of arrays to a 1D array ?

Please note that the actual size of C is 2 M * 4 and item_list is 20000.

Upvotes: 0

Views: 75

Answers (1)

Max von Hippel
Max von Hippel

Reputation: 2970

You want to flatten the list. For instance:

a = [[0, 1], [0, 1, 2], [1]]
flat_list = [item for sublist in a for item in sublist]


In [5]: flat_list

Out[5]: [0, 1, 0, 1, 2, 1]

In the case of your particular code, you could do:

rind = [[i for i in range(len(C)) if item in C[i]]
        for item in item_list]
rind = [item for sublist in rind for item in sublist]

Alternatively, you could do it in one line like this:

rind = list(map(set, [[i for i in range(len(C)) if item in C[i]]
        for item in item_list]))

Upvotes: 1

Related Questions