CODEWITHSUNDEEP
pythonarraysnumpy
junhuizh
junhuizh

Reputation: 412

Randomly pick an element in every N elements in numpy array

How to randomly pick an element for every N elements in a numpy 1D array?

For example, I have a numpy array [1,2,3,4,5,6,7,8,9,10] and I want to randomly pick one element in every two elements. One example is [1,4,5,7,10], which means 1 is randomly selected in [1,2], 4 is randomly selected from [3,4], and so on.

I tried to reshape and shuffle the array but I cannot find an elegant way to do it without a loop.

Upvotes: 2

Views: 490

Answers (1)

Divakar
Divakar

Reputation: 221564

For the length divisible by window/group length, we can reshape to a 2D array with the length as the number of cols and then select one random element per row -

def random_pick(a, W):
    b = a.reshape(-1,W)
    idx = np.random.randint(0,b.shape[1], len(b))
    return b[np.arange(len(idx)), idx]

Sample run -

In [17]: a
Out[17]: array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])

In [18]: np.random.seed(0)

In [19]: random_pick(a, W=2)
Out[19]: array([ 1,  4,  6,  7, 10])

Compact version

This also translates to a more compact version -

In [81]: W = 2

In [82]: np.random.seed(0)

In [83]: a[np.random.randint(0,W,len(a)//W) + np.arange(0,len(a),W)]
Out[83]: array([ 1,  4,  6,  7, 10])

Generic window length

To make it generic so that any window length could be fed, it would modify to -

def random_pick_generic(a, W):
    L = W*(len(a)//W)
    b = a[:L].reshape(-1,W)
    idx = np.random.randint(0,b.shape[1], len(b))
    out = b[np.arange(len(idx)), idx]
    if len(a[L:])>0:
        out = np.r_[out, np.random.choice(a[L:])]
    return out

Sample run -

In [50]: a
Out[50]: array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])

In [51]: np.random.seed(0)

In [52]: random_pick_generic(a, W=2)
Out[52]: array([ 1,  4,  6,  7, 10])

In [53]: np.random.seed(0)

In [54]: random_pick_generic(a, W=3)
Out[54]: array([ 1,  5,  7, 10])

Extend to 3D case : [batch_size, time_step, channel]

The solution assuming divisible window length random_pick would modify to -

b = a.reshape(-1,a.shape[1]//W,W,a.shape[2])
idx = np.random.randint(0,b.shape[2], b.shape[1])
out = b[:,np.arange(len(idx)), idx]

Upvotes: 1

Related Questions