Finding the distance to the nearest 0 in a list

Question:

I have a code that finds the distance of each element of the list to the nearest 0:

lst = [2, 1, 0, 3, 5]
length = len(lst)
dist = []

for i in lst:
    if lst.index(0) > lst.index(i):
        diff = lst.index(0) - lst.index(i)
        dist.append(diff)
    elif lst.index(0) < lst.index(i):
        diff = abs(lst.index(i)) - abs(lst.index(0))
        dist.append(diff)
    elif lst.index(0) == lst.index(i):
        dist.append(0)
print(dist)

# Output: [2, 1, 0, 1, 2]

Currently it works only when there is one 0 in a list.

How to modify code so that it works with any number of 0s in the list?

If second 0 is added to the list, than the output gets broken.

Asked By: nikitushu2

||

Answers:

index(0) will only get you the first occurrence of 0 — for that matter, index(i) will only get you the first occurrence of the value i, so if you had (say) multiple 5s in your list, your output would calculate the distance between the first 5 and the first 0, and report that for both 5s.

What I might suggest is getting the indices of all the zeroes in a list comprehension:

zeroes = [i for i, val in enumerate(lst) if val == 0]

and then you can get the distance to the closest zero for each index by taking the min distance across that zeroes list (this isn’t the absolute most efficient way, since it’s O(n^2) for a list that’s nothing but zeroes, but it has the advantage of being extremely simple):

dist = [min(abs(i - z) for z in zeroes) for i in range(len(lst))]

If we pop those two lines of code into a function it’s easy to try it out with different lists:

>>> def zero_dist(lst):
...     zeroes = [i for i, val in enumerate(lst) if val == 0]
...     return [min(abs(i - z) for z in zeroes) for i in range(len(lst))]
...
>>> zero_dist([2, 1, 0, 3, 5])
[2, 1, 0, 1, 2]
>>> zero_dist([5, 1, 0, 3, 0, 5])
[2, 1, 0, 1, 0, 1]

If you want something that’s O(n) even when there are lots of zeroes, you need to make sure that you’re not iterating over the entire original list of numbers or the list of all zeroes within your main loop body. Here’s one possible approach that’s based on building the list of zero indices and popping items off it as you go through the main loop, so that you’re only ever looking at the two nearest zero indices:

def zero_dist(nums: list[int]) -> list[int]:
    zeroes = [i for i, val in enumerate(nums) if val == 0][::-1]
    dist: list[int] = []
    z = next_z = zeroes.pop()  # raises IndexError if no zeroes
    for i in range(len(nums)):
        if i == next_z:
            z, next_z = next_z, zeroes.pop() if zeroes else -1
        dist.append(min(abs(i - z), abs(i - next_z)))
    return dist

Another approach is to build dist in two passes — first sweep forward through the list and decrease the potential value of dist based on how long it’s been since you last saw a zero, then sweep backward and do the same thing, setting the minimum value each time.

def zero_dist(nums: list[int]) -> list[int]:
    d = len(nums)
    dist = [d for _ in nums]  # will return this if no zeroes
    for e in (
        enumerate(nums),
        ((len(nums) - i, n) for i, n in enumerate(reversed(nums), 1))
    ):
        for i, n in e:
            if n == 0:
                d = 0
            else:
                d += 1
            dist[i] = min(d, dist[i])
    return dist
Answered By: Samwise

This works for any number of Zeros by first finding their locations and working out all distances from zeros then taking the minimums:

lst = [2, 1, 0, 3, 0, 5, 0, 7]

pos0 = [i for i,j in enumerate(lst) if j == 0]

result = []
for i in pos0:
    res = [abs(i-a) for a, b in enumerate(lst)]
    result.append(res)
    

final = [min(a) for a in zip(*result)]

print(final)

# prints [2, 1, 0, 1, 0, 1, 0, 1]
Answered By: user19077881

There is a way to implement a one-pass solution that will be O(n) time and O(1) space. This can even be written as a generator function that will accept any iterable as input.

For series or non-zeroes (except the first and last) the distances will increase from one up to half the size of the series and decrease back to 1. If the list starts with non-zeroes, their distances will decrease only. If the list ends with non-zeroes, their distance will only increase. So if we build the result based on increasing (forward) distances, we only need to generate the increasing and/or decreasing series of numbers based on the length of sequences of non-zero values.

For example:

A sequence of numbers with two zero entries contains all the possible patterns which we can break down as follows:

 9, 9, 9, 9, 0, 9, 9, 9, 9, 9, 0, 9, 9, 9, 9
|-----------|  |--------------|  |----------| 
|  leading  |  | intermediate |  | trailing |
|-----------|  |--------------|  |----------| 
 4, 3, 2, 1, 0, 1, 2, 3, 2, 1, 0, 1, 2, 3, 4 

The leading sequence of non-zero (up to the first zero) will always have decreasing distances

The intermediate sequence of non-zero (between zeros) will increase up to half its length and then decrease back to 1.

The trailing sequence of non-zero (after the last zero) will only increase.

Based on this the generator function can produce the result as it moves forward through the iterator, generating the distance numbers from ranges as it encounters the zero values.

def dist20(L):
    seenZero = False
    nonZeros = 0
    for n in L:
        if n:
            nonZeros += 1                             # length of non-zeros
        else:
            if seenZero:
                yield from range(1,(nonZeros+1)//2+1) # intermediate
                yield from range(nonZeros//2,0,-1)
            else:
                yield from range(nonZeros,0,-1)       # leading
            yield 0                                   # zero itself
            nonZeros = 0
            seenZero = True
    if seenZero:
        yield from range(1,nonZeros+1)                # trailing
    else:
        yield from (None for _ in range(nonZeros))    # no zeros found

output:

L = [2, 1, 0, 3, 5, 3, 0, 2, 1]
iL = iter(L)                     # doesn't have to be an actual iterator
print(*dist20(iL))               # but for demonstration purposes it is.
# 2 1 0 1 2 1 0 1 2              # print(*dist20(L)) would also work
Answered By: Alain T.
Categories: questions Tags: , , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.