Getting the row index of a specific value for each column using numpy

user_14388049 0
Viewed 1003

I have a matrix of 10,000 by 10,000 filled with 1s and 0s. What i want to do is to go through each column and find the rows that contain the value 1.

Then I want to store it in a new matrix with 2 columns : column 1 = column index and Column 2 = an array of row indices that contain 1. There are some columns that do not have any 1s at all, in which case it would be an empty array.

Trying to do a for loop again but it is computationally inefficient.

I tried with a smaller matrix

#sample matrix
n = 4
mat = [[randint(0,1) for _ in range(n)] for _ in range(n)]

arr = np.random.randint(0, size=(4, 2))

for col in range(n):
    arr[n][1] = n
    arr[n][2] = np.where(col == 1)

but this runs quite slowly for a 10,000 by 10,000 matrix. I am wondering if this is right and if there was a better way?

numpy python
1 Answers

Getting indices where a[i][j] == 1

You can get the data that you're looking for (locations of ones within a matrix of zeroes and ones) efficiently using numpy.argwhere() or numpy.nonzero(), however you will not be able to get them in the format specified in your original question using NumPy ndarrays alone.

You could achieve the data in you're specified format using a combination of ndarrays and standard Python lists, however since efficiency is paramount given the size of the data you are working with I would think it best to focus on getting the data rather than getting it in the format of an ndarray of irregular Python lists.

You can always reformat the results (indices of 1 within your matrix) following computation if the format you have mentioned is a hard requirement, and this way your code will benefit from optimisations provided by NumPy during the heavy computation - reducing the execution time of your procedure overall.

Example using np.argwhere()

import numpy as np

a = np.random.randint(0, 2, size=(4,4))
b = np.argwhere(a == 1)

print(f'a\n{a}')
print(f'b\n{b}')

Output

a
[[1 1 1 1]
 [0 0 0 0]
 [1 0 1 0]
 [1 1 1 1]]
b
[[0 0]
 [0 1]
 [0 2]
 [0 3]
 [2 0]
 [2 2]
 [3 0]
 [3 1]
 [3 2]
 [3 3]]

As you can see, np.argwhere(a == 1) returns an ndarray whose values are ndarrays containing the indices of locations in a whose values (x) meet the condition x == 1.

I gave the above method with a = np.random.randint(0, 2, size=(10000,10000) a try on my laptop (nothing fancy) a few times and it finished at around 3-5 seconds each time.

Getting row indices where all values != 1

If you want to store all row indices of a containing no values == 1, the most straightforward way (assuming you are using my example code above) would probably be by using numpy.setdiff1d() to return an array of row indices that are not present within b - i.e. the set difference between an array containing all row indices of a and the 1d array b[0] which will be row indices of all values in a that are != 1.

Assuming the same a and b as the above example.

c = np.setdiff1d(np.arange(a.shape[0]), b[:, 0])
print(c)

Output

array([1])

In the above example c = [1] as 1 is the only row index in a that doesn't contain any values == 1.

It is worth noting that if a is defined as np.random.randint(0, 2, size=(10000,10000), the probability of c being anything but a zero-length (i.e. empty) array is vanishingly small. This is because for a row to contain no values == 1, np.random would have to return 0 10,000 times in a row to fill a row with 0.

Why use multiple NumPy arrays?

I know that it may seem strange to use b and c to store results pertaining to locations where a == 1 and a != 1 respectively. Why not just use an irregular list as outlined in your original question?

The answer in short is efficiency. By using NumPy arrays you will be able to vectorise computations on your data and largely avoid costly Python loops, the benefits of which will be magnified considerably as reflected in time spent on execution given the size of the data you are working with.

You can always store your data in a different format that is more human friendly and map it back to NumPy as required, however the above examples will likely increase efficiency substantially at execution time when compared to the example in your original question.

user_120459070
Related