如何使用numpy将矢量填充和/或截断到指定的长度？

Question

我有几个列表:

a = [1,2,3]
b = [1,2,3,4,5,6]

它们的长度可变.

我想返回一个长度为5的向量,这样如果输入列表长度<5,那么它将在右边用零填充,如果它> 5,那么它将在第5个元素处被截断.

例如,输入a将返回np.array([1,2,3,0,0]),输入b将返回np.array([1,2,3,4,5]).

我觉得我应该可以使用np.pad,但我似乎无法遵循文档.

Answer 1

这可能是缓慢或快速,我不确定,但它适用于您的目的.

In [22]: pad = lambda a,i : a[0:i] if len(a) > i else a + [0] * (i-len(a))

In [23]: pad([1,2,3], 5)
Out[23]: [1, 2, 3, 0, 0]

In [24]: pad([1,2,3,4,5,6,7], 5)
Out[24]: [1, 2, 3, 4, 5]