Permutations over a subarray in python

Question

I have an array of identifiers that have been grouped into triples. For each group, I would like to randomly assign them to one of the three sets and store these assignments in another array. So, for a given array of grouped identifiers (I predict them):

groupings = array([1,1,1,2,2,2,3,3,3]) 

A possible way out would be

 assignments = array([0,1,2,1,0,2,2,0,1]) 

Ultimately, I would like to be able to generate many of these assignment lists and do it efficiently. My current method is to create an array of zeros and set each consecutive subframe of length 3 to a random permutation of 3.

 assignment = numpy.zeros((12,10),dtype=int) for i in range(0,12,3): for j in range(10): assignment[i:i+3,j] = numpy.random.permutation(3) 

Is there a better / faster way?

Answer 1

Two things I can think of:

• instead of visiting the 3 row * 1 column 2D array in the inner loop, try visiting it 1*3 . Typically, accessing a two-dimensional array horizontally is faster at first than vertically, as it gives you better spatial locality, which is good for caching.

• instead of running numpy.random.permutation(3) every time if 3 is fixed and is a small number, try creating arrays of permutations in advance and saving them into a constant array of the array, for example: (array([0,1,2]), array([0,2,1]), array([1,0,2])...) . You just need to randomly select one array from it each time.