Iterating elements of a numpy array by cycling

Viewed 199

I have a pretty straightforward question that I couldn't figure out quickly from the numpy reference documentation.

Say I have a numpy array labels = np.array([1, 2, 3]).
I have another array arr = np.array([1, 1, 1, 2, 2, 2, 3, 3, 3]).
I'd like to randomly sample indices of arr. Say our indices are [0, 3, 6].

Now, I'd like to make the elements corresponding to those indices cycle by one in labels. So, since arr[0] == 1, we would set arr[0] = 2. Since arr[3] == 2, we would set arr[3] = 3. Since arr[6] == 3, we would set arr[6] = 1.

so, to recap:

arr = np.array([1, 1, 1, 2, 2, 2, 3, 3, 3])
labels = np.unique(arr)
idx = np.array([0, 3, 6])  # randomly generate indices [0, 3, 6]
new_arr == np.array(2, 1, 1, 3, 2, 2, 1, 3, 3])  # this is the array I want

This seems like it'd be pretty straightforward, but I can't find an elegant way of doing this quickly!

1 Answers

If I understood you right, you might combine np.roll with np.random.randint, i.e.:

sample_idx = np.roll(np.random.randint(len(arr), 3), shift=1)                                                                                                                                                                                            
sample = [labels[i] for i in sample_idx]

That is, randomly sample 3 indexes and roll them by one. This should be the final indexes to be re-mapped to your labels.

EDIT

Got it. You might do it in two steps:

new_arr = arr.copy()
new_arr[idx] = np.roll(labels, -1)  # fixed the shift to -1

Output:

array([2, 1, 1, 3, 2, 2, 1, 3, 3])
Related