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

Question

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.

Max von Hippel · Answer 1 · 2018-06-01T22:09:20.367

1

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]))

edited Jun 01 '18 at 22:09

answered Jun 01 '18 at 22:03

Max von Hippel

1

I recommend using `a_set = list(map(set, a))` in your list comprehension. The one-off conversion cost is likely less than the benefit of O(1) lookup in loops. – jpp Jun 01 '18 at 22:05
1

Good idea. I'll write the map and edit accordingly. – Max von Hippel Jun 01 '18 at 22:06
(Just realized you cleverly use `set` for flattening here, so no actual map has to be written. Cool trick!) – Max von Hippel Jun 01 '18 at 22:09

1 Answers1