如何根据长度向量将numpy数组元素设置为零

时间:2017-09-11 18:12:52

标签: python numpy tensorflow

说我有以下数组d

>>> d = np.arange(25).reshape(5,5)
>>> d
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

我有向量l告诉我d中每行的长度:

>>> l = np.array([2,2,3,4,5])
>>> l
array([2, 2, 3, 4, 5])

如何将d行中长度超过l中指定的元素归零,以获得此结果:

>>> # zero out end of rows in `d`
>>> d
array([[ 0,  1,  0,  0,  0],
       [ 5,  6,  0,  0,  0],
       [10, 11, 12,  0,  0],
       [15, 16, 17, 18,  0],
       [20, 21, 22, 23, 24]])

这是针对张量流的,因此等效的tf会更好。谢谢!

1 个答案:

答案 0 :(得分:1)

使用broadcasting以矢量化方式创建尾随位置的蒙版,然后只需使用boolean-indexing重置输入数组中的那些,就像这样 -

d[l[:,None] <= np.arange(d.shape[1])] = 0

l2Dl[:,None] tensorflow的扩展名的等效值为:tf.expand_dims(l, 1)tf.expand_dims(l, -1)

示例运行 -

In [83]: d
Out[83]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

In [84]: l = np.array([2,2,3,4,5])

# mask of trailing places
In [85]: l[:,None] <= np.arange(d.shape[1])
Out[85]: 
array([[False, False,  True,  True,  True],
       [False, False,  True,  True,  True],
       [False, False, False,  True,  True],
       [False, False, False, False,  True],
       [False, False, False, False, False]], dtype=bool)

In [86]: d[l[:,None] <= np.arange(d.shape[1])] = 0

In [87]: d
Out[87]: 
array([[ 0,  1,  0,  0,  0],
       [ 5,  6,  0,  0,  0],
       [10, 11, 12,  0,  0],
       [15, 16, 17, 18,  0],
       [20, 21, 22, 23, 24]])