Efficiently index 2d numpy array using two 1d arrays

Question

I have a large 2d numpy array and two 1d arrays that represent x/y indexes within the 2d array. I want to use these 1d arrays to perform an operation on the 2d array. I can do this with a for loop, but it's very slow when working on a large array. Is there a faster way? I tried using the 1d arrays simply as indexes but that didn't work. See this example:

import numpy as np

# Two example 2d arrays
cnt_a   =   np.zeros((4,4))
cnt_b   =   np.zeros((4,4))

# 1d arrays holding x and y indices
xpos    =   [0,0,1,2,1,2,1,0,0,0,0,1,1,1,2,2,3]
ypos    =   [3,2,1,1,3,0,1,0,0,1,2,1,2,3,3,2,0]

# This method works, but is very slow for a large array
for i in range(0,len(xpos)):
    cnt_a[xpos[i],ypos[i]] = cnt_a[xpos[i],ypos[i]] + 1

# This method is fast, but gives incorrect answer
cnt_b[xpos,ypos] = cnt_b[xpos,ypos]+1


# Print the results
print 'Good:'
print cnt_a
print ''
print 'Bad:'
print cnt_b

The output from this is:

Good:
[[ 2.  1.  2.  1.]
 [ 0.  3.  1.  2.]
 [ 1.  1.  1.  1.]
 [ 1.  0.  0.  0.]]

Bad:
[[ 1.  1.  1.  1.]
 [ 0.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  0.  0.  0.]]

For the cnt_b array numpy is obviously not summing correctly, but I'm unsure how to fix this without resorting to the (v. inefficient) for loop used to calculate cnt_a.

Answer 1

Another approach by using 1D indexing (suggested by @Shai) extended to answer the actual question:

>>> out = np.zeros((4, 4))
>>> idx = np.ravel_multi_index((xpos, ypos), out.shape) # extract 1D indexes
>>> x = np.bincount(idx, minlength=out.size)
>>> out.flat += x

np.bincount calculates how many times each of the index is present in the xpos, ypos and stores them in x.

Or, as suggested by @Divakar:

>>> out.flat += np.bincount(idx, minlength=out.size)

Answer 2

We could compute the linear indices, then accumulate into zeros-initialized output array with np.add.at. Thus, with xpos and ypos as arrays, here's one implementation -

m,n = xpos.max()+1, ypos.max()+1
out = np.zeros((m,n),dtype=int)
np.add.at(out.ravel(), xpos*n+ypos, 1)

Sample run -

In [95]: # 1d arrays holding x and y indices
    ...: xpos    =   np.array([0,0,1,2,1,2,1,0,0,0,0,1,1,1,2,2,3])
    ...: ypos    =   np.array([3,2,1,1,3,0,1,0,0,1,2,1,2,3,3,2,0])
    ...: 

In [96]: cnt_a   =   np.zeros((4,4))

In [97]: # This method works, but is very slow for a large array
    ...: for i in range(0,len(xpos)):
    ...:     cnt_a[xpos[i],ypos[i]] = cnt_a[xpos[i],ypos[i]] + 1
    ...:     

In [98]: m,n = xpos.max()+1, ypos.max()+1
    ...: out = np.zeros((m,n),dtype=int)
    ...: np.add.at(out.ravel(), xpos*n+ypos, 1)
    ...: 

In [99]: cnt_a
Out[99]: 
array([[ 2.,  1.,  2.,  1.],
       [ 0.,  3.,  1.,  2.],
       [ 1.,  1.,  1.,  1.],
       [ 1.,  0.,  0.,  0.]])

In [100]: out
Out[100]: 
array([[2, 1, 2, 1],
       [0, 3, 1, 2],
       [1, 1, 1, 1],
       [1, 0, 0, 0]])

Answer 3

you can iterate on both lists at once, and increment for each couple (if you are not used to it, zip can combine lists)

for x, y in zip(xpos, ypos):
    cnt_b[x][y] += 1

But this will be about the same speed as your solution A. If your lists xpos/ypos are of length n, I don't see how you can update your matrix in less than o(n) since you'll have to check each pair one way or an other.

Other solution: you could count (with collections.Counter possibly) the similar index pairs (ex: (0, 3) etc...) and update the matrix with the count value. But I doubt it would be much faster, since you the time gained on updating the matrix would be lost on counting multiple occurrences.

Maybe I am totally wrong tho, in which case I'd be curious too to see a not o(n) answer

Answer 4

I think you are looking for ravel_multi_index funciton

lidx = np.ravel_multi_index((xpos, ypos), cnt_a.shape)

converts to "flatten" 1D indices into cnt_a and cnt_b:

np.add.at( cnt_b, lidx, 1 )