Distances between same numbers in list using single for loop

Question

I have an assignment to get distances of every occurrence of same number in string and its time complexity should be O(n) so it shouldn't use nested for loops.

For example if my string contains "100101" and I need to get distances between ones it's total distance would be 10. (Since first and second has distance of 3, first and last has 5 and second and last has 2).

I do get the correct answers using nested for loops but I don't understand how this should be implemented without nested loops.

My current code:

def pairs(s):
    array = []
    total = 0

    for i in range(len(s)):
        if s[i] == "1":
            array.append(i)

    for k in reversed(array):
        array.remove(k)
        for j in reversed(array):
            total += k - j
        
    return total



if __name__ == "__main__":
    print(pairs("100101")) # 10
    print(pairs("101")) # 2
    print(pairs("100100111001")) # 71

Answer 1

This can be indeed computed in linear time (I adapted the classic greedy algorithm used to compute the sum of Manhattan distances of sorted points):

def pairs(numbers, target):
    result = s = i = 0
    for j, n in enumerate(numbers):
        if n == target:
            result += (j * i - s)
            s += j
            i += 1
    return result

Some examples:

>>> pairs('100101', '1')
10
>>> pairs('100101', '0')
6
>>> pairs('100010101', '1')
26

Answer 2

Another variant:

def pairs(s):
    total = ones = distance = 0
    for digit in s:
        if digit == '1':
            ones += 1
            total += distance
        distance += ones
    return total

ones: The number of 1s seen so far.

total: The sum of distances between the 1s seen so far.

distance: the sum of distances of previous 1s to the current position.

When you come across another 1, count it (ones += 1) and add its sum of distances to all previous 1s to the overall total (total += distance).

When you move to the next position, you move one position away from all 1s seen so far, hence distance += ones.