在 python 中按组随机播放数组

Question

Let's assume I have two arrays:

values = [1,2,3,4,5,6,7,8,9]
groups = [0,0,0,1,1,2,2,3,4]

Is it possible to shuffle "values" only within groups? Eg elements in group 0 (1,2,3) are going to be shuffled only with each other, elements in group 1 (4,5) are going to be shuffled with each other and so on.

I have huge numpy arrays, is there any efficient way to do so?

Answer 1

You can do it this way:

import numpy as np
np.random.seed(133)

values = np.array([1,2,3,4,5,6,7,8,9])
groups = np.array([0,0,0,1,1,2,2,3,4])

for index in np.unique(groups):
    mask = groups==index
    values[mask] = np.random.permutation(values[mask])

print(values)

Output:

[3 1 2 5 4 6 7 8 9]

Answer 2

Assuming that your group numbers are always in ascending order, you can leverage the fact that Python's sort is stable to shuffle the values/groups as a whole and then sort the result only by groups. Combine the group numbers and values into a single list of tuples that you shuffle. Then sort that list of tuple using only the group as sort key and extract only the value part

values = [1,2,3,4,5,6,7,8,9]
groups = [0,0,0,1,1,2,2,3,4]

import random

shuffled = random.sample([*zip(groups,values)],len(values))
values   = [v for g,v in sorted(shuffled,key=lambda gs:gs[0])]

print(values)
print(groups)

[3, 1, 2, 5, 4, 7, 6, 8, 9]
[0, 0, 0, 1, 1, 2, 2, 3, 4]

If your group identifiers are not ordered (or not consecutive), you will need to form groups (of indexes), shuffle them group by group and place the shuffled values at the specific subset of positions corresponding to each group:

values = [1,2,3,4,5,6,7,8,9]
groups = [0,1,0,1,0,2,3,3,2]

import random

gIndex = dict() # grouping dictionary {groupId:[indexes]}
for i,g in enumerate(groups): 
    gIndex.setdefault(g,[]).append(i) # value indexes by group id
shuffled = [None]*len(groups)         # resulting shuffled value list
for indexes in gIndex.values():       # shuffle indexes by group
    for i,j in zip(indexes,random.sample(indexes,len(indexes))):
        shuffled[i] = values[j]       # map old positions to new position
        
print(values)
print(groups)
print(shuffled)

[1, 2, 3, 4, 5, 6, 7, 8, 9] # original order
[0, 1, 0, 1, 0, 2, 3, 3, 2] # group identifiers
[3, 4, 1, 2, 5, 9, 8, 7, 6] # shuffled order (within groups)

You could also do ti using the first sorting technique but it wouldn't be as efficient as using a dictionary:

indexes    = sorted((g,i) for i,g in enumerate(groups))
newIndexes = sorted(random.sample(indexes,len(indexes)),key=lambda gi:gi[0])
shuffled   = [None]*len(values)
for (_,i),(_,j) in zip(indexes,newIndexes):
    shuffled[i] = values[j]