Padding or truncating a Python list

Question

I'd like to truncate or pad a list. E.g. for size 4:

[1,2,3] -> [1,2,3,0]
[1,2,3,4,5] -> [1,2,3,4]

I can see a couple of ways:

def trp(l, n):
    """ Truncate or pad a list """
    r = l[:n]
    if len(r) < n:
        r.extend([0] * (n - len(r)))
    return r

Or a shorter, but less efficient:

map(lambda x, y: x if x else 0, m[0:n], [0] * n)

Is there a more elegant way of doing this?

Answer 1

You can use itertools module to make it completely lazy, like this

>>> from itertools import repeat, chain, islice
>>> def trimmer(seq, size, filler=0):
...     return islice(chain(seq, repeat(filler)), size)
... 
>>> list(trimmer([1, 2, 3], 4))
[1, 2, 3, 0]
>>> list(trimmer([1, 2, 3, 4, 5], 4))
[1, 2, 3, 4]

Here, we chain the actual sequence with the infinite repeater with the filler value. And then we slice the chained iterator to size.

So, when if the sequence has lesser number of elements than size, chain will start consuming the repeat. If the sequence has at least size elements, then chain will not even have to use the repeat.

The main advantage of this method is that, the complete trimmed or padded list is not created in memory, unless asked for it. So, if all you are going to do is to iterate it, then you can simply iterate it like this

>>> for item in trimmer([1, 2, 3, 4, 5], 4):
...     print(item * 2)
...     
... 
2
4
6
8

Or, if you want to use it with another trimmed or padded list, then you can still do that without creating an actual list, like this

>>> for item in chain(trimmer([1, 2, 3], 4), trimmer([1, 2, 3, 4, 5], 4)):
...     print(item, item * 2)
...     
... 
1 2
2 4
3 6
0 0
1 2
2 4
3 6
4 8

Laziness Rocks ;-)

Answer 2

Slicing using an index greater than the length of a list just returns the entire list.

Multiplying a list by a negative value returns an empty list.

That means the function can be written as:

def trp(l, n):
    return l[:n] + [0]*(n-len(l))

trp([], 4)
[0, 0, 0, 0]

trp([1,2,3,4], 4)
[1, 2, 3, 4]

trp([1,2,3,4,5], 4)
[1, 2, 3, 4]

trp([1,2,3], 4)
[1, 2, 3, 0]

---

In [1]: a = [1,2,3]

In [2]: a[:4]
Out[2]: [1, 2, 3]

In [3]: [0]*0
Out[3]: []

In [4]: [0]*-1
Out[4]: []