Numpy collapse columns according to list
<p>In numpy, I have a d x n array A and a list L of length n describing where I want each column of A to end up in matrix B. The idea is that column i of matrix B is the sum of all columns of A for which the corresponding value in L is i. I can do this with a loop:</p>

<pre><code>A = np.arange(15).reshape(3,5)
L = [0,1,2,1,1]
n_cols = 3
B = np.zeros((len(A), n_cols)) 
# assume I know the desired number of columns, 
# which is also one more than the maximum value of L
for i, a in enumerate(A.T):
    B[:, L[i]] += a
</code></pre>

<p>Is there a way to do this by slicing array A (or in some otherwise vectorized manner)?</p>

<p>You are sum-reducing/collapsing the columns off <code>A</code> using <code>L</code> for selecting those columns. Also, you are updating the columns of output array based on the uniqueness of <code>L</code> elems. </p>

<p>Thus, you can use <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.ufunc.reduceat.html" rel="nofollow"><code>np.add.reduceat</code></a> for a vectorized solution, like so -</p>

<pre><code>sidx = L.argsort()
col_idx, grp_start_idx = np.unique(L[sidx],return_index=True)
B_out = np.zeros((len(A), n_cols))
B_out[:,col_idx] = np.add.reduceat(A[:,sidx],grp_start_idx,axis=1)
</code></pre>

<hr>

<p>Runtime test -</p>

<pre><code>In [129]: def org_app(A,n_cols):
     ...:     B = np.zeros((len(A), n_cols)) 
     ...:     for i, a in enumerate(A.T):
     ...:         B[:, L[i]] += a
     ...:     return B
     ...: 
     ...: def vectorized_app(A,n_cols):
     ...:     sidx = L.argsort()
     ...:     col_idx, grp_start_idx = np.unique(L[sidx],return_index=True)
     ...:     B_out = np.zeros((len(A), n_cols))
     ...:     B_out[:,col_idx] = np.add.reduceat(A[:,sidx],grp_start_idx,axis=1)
     ...:     return B_out
     ...: 

In [130]: # Setup inputs with an appreciable no. of cols &amp; lesser rows
     ...: # so as that memory bandwidth to work with huge number of 
     ...: # row elems doesn't become the bottleneck
     ...: d,n_cols = 10,5000
     ...: A = np.random.rand(d,n_cols)
     ...: L = np.random.randint(0,n_cols,(n_cols,))
     ...: 

In [131]: np.allclose(org_app(A,n_cols),vectorized_app(A,n_cols))
Out[131]: True

In [132]: %timeit org_app(A,n_cols)
10 loops, best of 3: 33.3 ms per loop

In [133]: %timeit vectorized_app(A,n_cols)
100 loops, best of 3: 1.87 ms per loop
</code></pre>

<p>As the number of rows become comparable with the number of cols in <code>A</code>, the high memory bandwidth requirements for the vectorized approach would offset any noticeable speedup from it.</p>

python numpy vectorization numpy-broadcasting