23

My goal is to take a list of unknown number of elements and extend/slice it to exactly n elements, padding lists that are too short with 0 and slicing lists that are too long.

For example

n = 10
foo = [1,2,3,4]
print some_func(foo,n) 

should return [1,2,3,4,0,0,0,0,0,0], and

n = 10
foo = [1,2,3,4,5,6,7,8,9,10,11,12]
print some_func(foo,n)

should return [1,2,3,4,5,6,7,8,9,10]


Right now I'm doing this:

def some_function(l, n):
    l.extend([0] * n)
    l = l[:n]
    return l

But that seems inefficient. Is there a more pythonic way of doing this?


EDIT: point of clarification, I am not trying to modify the original array, I am returning a new array that can be a shallow copy.

Hans Z
  • 4,664
  • 2
  • 27
  • 50

6 Answers6

18

The only potentially "more Pythonic" way to do this is pretty much the way you have it, but skip over the extra variable allocation and just return the result directly. You could also definitely make it more Pythonic by conforming to PEP8 with your function and variable names.

Besides that, you can improve your efficiency by adding only as many zeroes as you need, rather than building a too-long list and then trimming it.

def pad_or_truncate(some_list, target_len):
    return some_list[:target_len] + [0]*(target_len - len(some_list))

Breaking it down, there are two cases represented here (discounting the trivial case where the input is exactly the right length already). Either the list is too long, or the list is too short.

If the list is too long, we just slice it. some_list[:target_len] takes care of that. Because the slice operation is friendly, this won't blow up if the target length is beyond the actual length of the list.

If the list is too short, we pad it with zeroes. I chose to do this by multiplying a list literal1, but you can use a list comprehension the exact same way2 if that's more your cup of tea. Just determine how many zeroes to add (either zero, if the list isn't too short, or target_len - len(some_list)), and concatenate a list composed of that many zeroes. Done and done!

If you want to make it an in-place operation (as your original example appears to be trying but failing to achieve; see @chepner's comment), you would just change return <...> to some_list[:] = <...>.


1Some brief timeit results indicated that literal multiplication is a bit quicker than the double-iteration implied by a list comprehension.

2For posterity, the list comprehension version would look something like:

return some_list[:target_len] + [0 for _ in range(target_len - len(some_list))]
Hans Z
  • 4,664
  • 2
  • 27
  • 50
Henry Keiter
  • 16,863
  • 7
  • 51
  • 80
8

It is quite inefficient to build the list bigger than you need it (particularly if n gets large); instead, only add the padding you need. Also, it is Python convention to return None if a function changes its argument(s) in-place:

def some_func(l, n, pad=0):
    if len(l) >= n:
        del l[n:]
    else:
        l.extend([pad] * (n - len(l)))

Example:

>>> l = [1, 2, 3, 4, 5]
>>> some_func(l, 3)
>>> l
[1, 2, 3]
>>> some_func(l, 5)
>>> l
[1, 2, 3, 0, 0]

Alternatively, return a new list:

def some_func(l, n, pad=0):
    if len(l) >= n:
        return l[:n]
    return l + ([pad] * (n - len(l)))
jonrsharpe
  • 115,751
  • 26
  • 228
  • 437
3

How about islice-ing the concatenation of the original list with a padding generator?

from itertools import islice, repeat
def fit_to_size(l, n):
    return list(islice(
        ( x
          for itr in (l, repeat(0))
          for x in itr ),
        n))

You might prefer this slightly more explicit implementation:

def fit_to_size(l, n):
    def gen():
        yield from l
        while True: yield 0
    return list(islice(gen(), n))
shx2
  • 61,779
  • 13
  • 130
  • 153
1

Why not use conditional logic?

def pad_or_slice(L, n):
    if len(L) < n:
        return L + ([0] * (n - len(L)))
    else:
        return L[:n]

This way you're only doing one of the two, not both, at the cost of checking the length of the list, which shouldn't be too costly.

Dave Yarwood
  • 2,866
  • 1
  • 17
  • 29
0

It's not significantly more efficient, but I'd do l = (l + [0] * (n - len(l)))[:n]. So your some_function would look like this

def some_function(list, n):
    return (list + [0] * (n - len(list)))[:n]
Tyler
  • 1,818
  • 2
  • 13
  • 22
0

In order to properly modify the original list, you'll need to delete a trailing slice from the object.

def some_function(l, n):
    l.extend([0] * n)
    del l[n:]
chepner
  • 497,756
  • 71
  • 530
  • 681