How can I pad and/or truncate a vector to a specified length using numpy?

Question

I have couple of lists:

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

which are of variable length.

I want to return a vector of length five, such that if the input list length is < 5 then it will be padded with zeros on the right, and if it is > 5, then it will be truncated at the 5th element.

For example, input a would return np.array([1,2,3,0,0]), and input b would return np.array([1,2,3,4,5]).

I feel like I ought to be able to use np.pad, but I can't seem to follow the documentation.

Answer 1

This might be slow or fast, I am not sure, however it works for your purpose.

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]