python - 广播？For 循环 --> numpy

Question

我有的：

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]])   

我想要做的是仅使用 numpy 或 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')]

我认为答案只是扩展以下行：

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

但我不确定那是什么。

谁能给我一些建议？提前致谢！

标签： pythonnumpy

Answer 1

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]])

您可以使用

列表理解

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

如果您希望在没有for循环的情况下执行此操作，则可以一起使用多个 numpy 函数。我会这样做：

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)))

上面的输出rnd_ar将是：

[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')]