Numpy: How to quickly replace equal values in an matrix?

user_2913106 0
Viewed 86

Lets say we have a rank 2 array a with n entries that contain integer values in {0,1,2,...,m}. Now for each of those integers I want to find the indices of the entries of a with this value (called index_i, index_j in the following examples). (So what I'm looking for is like np.unique(...,return_index=True) but for 2d arrays and with the possibility to return all indices of each unique value.)

A naive approach would involve using boolean indexing which would result in O(m*n) operations (see below), but I'd like to only have O(n) operations. While I found a solution to do that, I feel like there should be a built in method or at least something that simplifies this - or that would at least remove these ugly loops:

import numpy as np
a = np.array([[0,0,1],[0,2,1],[2,2,1]])
m = a.max()


#"naive" in O(n*m)
i,j = np.mgrid[range(a.shape[0]), range(a.shape[1])]
index_i = [[] for _ in range(m+1)]
index_j = [[] for _ in range(m+1)]
for k in range(m+1):
  index_i[k] = i[a==k]
  index_j[k] = j[a==k]

#all the zeros:
print(a[index_i[0], index_j[0]])
#all the ones:
print(a[index_i[1], index_j[1]])
#all the twos:
print(a[index_i[2], index_j[2]])


#"sophisticated" in O(n)

index_i = [[] for _ in range(m+1)]
index_j = [[] for _ in range(m+1)]
for i in range(a.shape[0]):
  for j in range(a.shape[1]):
    index_i[a[i,j]].append(i)
    index_j[a[i,j]].append(j)

#all the zeros:
print(a[index_i[0], index_j[0]])
#all the ones:
print(a[index_i[1], index_j[1]])
#all the twos:
print(a[index_i[2], index_j[2]])

Try it online!

(Note that I will need these indices for write access later, that is, for replacing the values stored in the array. But between these operations I do need to the 2d structure.)

indexing numpy arrays python
1 Answers

Here's one based on sorting with the intention of having minimal work when iterating to save as a dictionary that has keys being the unique elements and the values as the indices -

shp = a.shape
idx = a.ravel().argsort()
idx_sorted = np.c_[np.unravel_index(idx,shp)]
count = np.bincount(a.ravel())
valid_idx = np.flatnonzero(count!=0)
cs = np.r_[0,count[valid_idx].cumsum()]
out = {e:idx_sorted[i:j] for (e,i,j) in zip(valid_idx,cs[:-1],cs[1:])}

Sample input, output -

In [155]: a
Out[155]: 
array([[0, 2, 6],
       [0, 2, 6],
       [2, 2, 1]])

In [156]: out
Out[156]: 
{0: array([[0, 0],
        [1, 0]]), 1: array([[2, 2]]), 2: array([[0, 1],
        [1, 1],
        [2, 0],
        [2, 1]]), 6: array([[0, 2],
        [1, 2]])}

If all integers in the sequence are covered in the array, we could simplify it a bit -

shp = a.shape
idx = a.ravel().argsort()
idx_sorted = np.c_[np.unravel_index(idx,shp)]
cs = np.r_[0,np.bincount(a.ravel()).cumsum()]
out = {iterID:idx_sorted[i:j] for iterID,(i,j) in enumerate(zip(cs[:-1],cs[1:]))}
user_32938810
Related