Search code examples
pythonalgorithmtime-complexitycomputer-science

Distances between same numbers in list using single for loop


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

Solution

  • 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