Quickly compute a randomized 3D numpy array from a 2D numpy array

Question

I have a two-dimensional array of integers, we will call it "A".

I want to create a three-dimensional array "B" of all 1s and 0s so that:

• for any fixed (i, j) sum(B[i,j,:])==A[ij] B[i,j,:] sum(B[i,j,:])==A[ij] , that is, B[i,j,:] contains A[i,j] 1s in it
• 1s are placed randomly in the 3rd dimension.

I know how to do this using standard python indexing, but this proves to be very slow.

I'm looking for a way to do this, taking advantage of features that Numpy can quickly do.

Here's how I would do it using standard indexing:

 B=np.zeros((X,Y,Z)) indexoptions=range(Z) for i in xrange(Y): for j in xrange(X): replacedindices=np.random.choice(indexoptions,size=A[i,j],replace=False) B[i,j,[replacedindices]]=1 

Can someone explain how I can do this faster?

Edit: Here is an example of "A":

 A=np.array([[0,1,2,3,4],[0,1,2,3,4],[0,1,2,3,4],[0,1,2,3,4],[0,1,2,3,4]]) 

in this case X = Y = 5 and Z> = 5

Answer 1

Essentially the same idea as @JohnZwinck and @DSM, but with a shuffle function to shuffle a given axis:

 import numpy as np def shuffle(a, axis=-1): """ Shuffle `a` in-place along the given axis. Apply numpy.random.shuffle to the given axis of `a`. Each one-dimensional slice is shuffled independently. """ b = a.swapaxes(axis,-1) # Shuffle `b` in-place along the last axis. `b` is a view of `a`, # so `a` is shuffled in place, too. shp = b.shape[:-1] for ndx in np.ndindex(shp): np.random.shuffle(b[ndx]) return def random_bits(a, n): b = (a[..., np.newaxis] > np.arange(n)).astype(int) shuffle(b) return b if __name__ == "__main__": np.random.seed(12345) A = np.random.randint(0, 5, size=(3,4)) Z = 6 B = random_bits(A, Z) print "A:" print A print "B:" print B 

Output:

 A: [[2 1 4 1] [2 1 1 3] [1 3 0 2]] B: [[[1 0 0 0 0 1] [0 1 0 0 0 0] [0 1 1 1 1 0] [0 0 0 1 0 0]] [[0 1 0 1 0 0] [0 0 0 1 0 0] [0 0 1 0 0 0] [1 0 1 0 1 0]] [[0 0 0 0 0 1] [0 0 1 1 1 0] [0 0 0 0 0 0] [0 0 1 0 1 0]]]