CODEWITHSUNDEEP
pythonnumpy
Kevin Choi
Kevin Choi

Reputation: 765

Broadcasting? For loop --> numpy

What I have:

rnd_ar = np.random.randint(5, size = (10,2))
rnd_ar
array([[4, 2],
       [2, 1],
       [2, 0],
       [3, 4],
       [0, 1],
       [0, 2],
       [1, 0],
       [0, 0],
       [3, 0],
       [1, 0]])

What I want to do is to simplify the following code by only using either numpy or pandas:

test = []
for i in rnd_ar:
    test.append(np.repeat(['O', 'T'],i))

test

[array(['O', 'O', 'O', 'O', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'T'], dtype='<U1'),
 array(['O', 'O'], dtype='<U1'),
 array(['O', 'O', 'O', 'T', 'T', 'T', 'T'], dtype='<U1'),
 array(['T'], dtype='<U1'),
 array(['T', 'T'], dtype='<U1'),
 array(['O'], dtype='<U1'),
 array([], dtype='<U1'),
 array(['O', 'O', 'O'], dtype='<U1'),
 array(['O'], dtype='<U1')]

I think the answer would be just expanding the following line:

np.repeat(np.array(['O', 'T']), rnd_ar[0])

but I am not sure what that is.

Can anyone give me some advice? Thanks in advance!

Upvotes: 0

Views: 95

Answers (1)

Maryam Bahrami
Maryam Bahrami

Reputation: 1104

import numpy as np
rnd_ar = np.random.randint(5, size = (10,2))
rnd_ar
array([[4, 3],
       [4, 1],
       [1, 0],
       [2, 4],
       [4, 2],
       [3, 4],
       [3, 2],
       [2, 1],
       [4, 3],
       [2, 1]])

You can use

list comprehension

test = [np.repeat(['O', 'T'],i) for i in rnd_ar]

If you wish to do it without a for loop, then you can use several numpy function all together. I will do it like this:

np.split(
    np.repeat(np.array(['O', 'T'] * rnd_ar.shape[0]).reshape(rnd_ar.shape),
              rnd_ar.flatten()),
    np.cumsum(np.sum(rnd_ar, axis=1)))

The output for the above rnd_ar will be:

[array(['O', 'O', 'O', 'O', 'T', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'O', 'O', 'T'], dtype='<U1'),
 array(['O'], dtype='<U1'),
 array(['O', 'O', 'T', 'T', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'O', 'O', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'O', 'T', 'T', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'O', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'T'], dtype='<U1'),
 array(['O', 'O', 'O', 'O', 'T', 'T', 'T'], dtype='<U1'),
 array(['O', 'O', 'T'], dtype='<U1'),
 array([], dtype='<U1')]

Upvotes: 1

Related Questions