Push all zeros to tail of int array

Question

My script pushes zeros to end of an int array. I'd like to be reviewed on efficiency, style, and obviously if there is a bug I would like to know.

def nonzero(a):
  ''' given an array of ints, push all zeros to the end '''
  zeros = [0 for i in range(a.count(0))]
  x = [ i for i in a if i != 0]
  x.extend(zeros)
  return(x)

Answer 1

1. The function does not preserve object identity. Zeros of other types get replaced with an integer:

   >>> nonzero([0, 0.0, 0j])
   [0, 0, 0]
   

   If the input only ever contains integers, then this isn't a problem, but it's worth remembering that Python has several different types of number.

2. The docstring is not quite right. "push all the zeros to the end" implies that the input array is modified, but that's not true: a new array is created and returned.

3. There's a neat trick for implementing this:

   def nonzero(a):
       """Return a copy of the iterable a with all zeros at the end."""
       return sorted(a, key=lambda x: x == 0)
   

   This works because False compares less than True, and because Python's sort is guaranteed to be stable: it never changes the relative order of elements with the same sort key.

Answer 2

nonzero is not a great name for this function. It implies a check to see if something is nonzero, not a simple sort method. It is hard to name something like this, but even the overly long zeroes_to_end would be clearer as it's not ambiguously implying something else. Likewise a and x aren't great, instead maybe source and result? (assuming you don't decide to modify the original list as Gareth Rees suggested)

Also instead of:

zeros = [0 for i in range(a.count(0))]

It'd be easier to just do this:

zeros = [0] * a.count(0)

which will give you the same thing.

Lastly you don't need parentheses around your return value. You can just do return a which reads as more Pythonic.

Answer 3

I'd improve your second comprehension: [i for i in a if i]. It's more Pythonic to write if i than if i!=0

This works because

a) Python makes some implicit type conversions and

b) in Python almost anything can be converted to boolean (e.g. when you type if some_var:). You will get false if some_var is not initialized, if it is 0 (not '0'), empty string, empty collection, etc.

Answer 4

If it is acceptable to work with numpy arrays. I prefer its syntax over lambda functions and list comprehensions.

def shift_zeros(a):
    return np.concatenate([a[a != 0], a[a == 0]])