Count number of tails since the last head

Question

Consider a sequence of coin tosses: 1, 0, 0, 1, 0, 1 where tail = 0 and head = 1.

The desired output is the sequence: 0, 1, 2, 0, 1, 0

Each element of the output sequence counts the number of tails since the last head.

I have tried a naive method:

def timer(seq):
    if seq[0] == 1: time = [0]
    if seq[0] == 0: time = [1]
    for x in seq[1:]:
        if x == 0: time.append(time[-1] + 1)
        if x == 1: time.append(0)
    return time

Question: Is there a better method?

Answer 1

Using NumPy:

import numpy as np 
seq = np.array([1,0,0,1,0,1,0,0,0,0,1,0])
arr = np.arange(len(seq))
result = arr - np.maximum.accumulate(arr * seq)
print(result)

yields

[0 1 2 0 1 0 1 2 3 4 0 1]

---

Why arr - np.maximum.accumulate(arr * seq)? The desired output seemed related to a simple progression of integers:

arr = np.arange(len(seq))

So the natural question is, if seq = np.array([1, 0, 0, 1, 0, 1]) and the expected result is expected = np.array([0, 1, 2, 0, 1, 0]), then what value of x makes

arr + x = expected

Since

In [220]: expected - arr
Out[220]: array([ 0,  0,  0, -3, -3, -5])

it looks like x should be the cumulative max of arr * seq:

In [234]: arr * seq
Out[234]: array([0, 0, 0, 3, 0, 5])

In [235]: np.maximum.accumulate(arr * seq)
Out[235]: array([0, 0, 0, 3, 3, 5])

Answer 2

Step 1: Invert l:

In [311]: l = [1, 0, 0, 1, 0, 1]

In [312]: out = [int(not i) for i in l]; out
Out[312]: [0, 1, 1, 0, 1, 0]

Step 2: List comp; add previous value to current value if current value is 1.

In [319]: [out[0]] + [x + y if y else y for x, y in zip(out[:-1], out[1:])]
Out[319]: [0, 1, 2, 0, 1, 0]

This gets rid of windy ifs by zipping adjacent elements.

Answer 3

Using itertools.accumulate:

>>> a = [1, 0, 0, 1, 0, 1]
>>> b = [1 - x for x in a]
>>> list(accumulate(b, lambda total,e: total+1 if e==1 else 0))
[0, 1, 2, 0, 1, 0]

accumulate is only defined in Python 3. There's the equivalent Python code in the above documentation, though, if you want to use it in Python 2.

It's required to invert a because the first element returned by accumulate is the first list element, independently from the accumulator function:

>>> list(accumulate(a, lambda total,e: 0))
[1, 0, 0, 0, 0, 0]

Answer 4

The required output is an array with the same length as the input and none of the values are equal to the input. Therefore, the algorithm must be at least O(n) to form the new output array. Furthermore for this specific problem, you would also need to scan all the values for the input array. All these operations are O(n) and it will not get any more efficient. Constants may differ but your method is already in O(n) and will not go any lower.