CSES - Distances

Your task is to implement a class that maintains a list of numbers and has a method for efficiently computing the sum of distances between all pairs of occurrences of a given number.

For example, when the list is [1,2,1,3,3,1,2,1][1,2,1,3,3,1,2,1], the distances for the number 11 are:

  • positions 00 and 22: distance 22
  • positions 00 and 55: distance 55
  • positions 00 and 77: distance 77
  • positions 22 and 55: distance 33
  • positions 22 and 77: distance 55
  • positions 55 and 77: distance 22

Thus the sum of the distances is 2+5+7+3+5+2=242+5+7+3+5+2=24.

In a file distances.py, implement the class DistanceTracker with the methods:

  • append(number): add number to the end of the list
  • sum(number): return the sum of distances for number

Both methods should work in O(1)O(1) time.

class DistanceTracker:
    def __init__(self):
        # TODO

    def append(self, number):
        # TODO

    def sum(self, number):
        # TODO

if __name__ == "__main__":
    tracker = DistanceTracker()

    tracker.append(1)
    tracker.append(2)
    tracker.append(1)
    tracker.append(3)
    tracker.append(3)
    tracker.append(1)
    tracker.append(2)
    tracker.append(1)

    print(tracker.sum(1)) # 24
    print(tracker.sum(2)) # 5
    print(tracker.sum(3)) # 1

    tracker.append(1)
    tracker.append(2)
    tracker.append(3)

    print(tracker.sum(1)) # 42
    print(tracker.sum(2)) # 16
    print(tracker.sum(3)) # 14

If a number does not appear on the list or appears only once, the desired result is 00:

    tracker = DistanceTracker()
    print(tracker.sum(1)) # 0
    tracker.append(1)
    print(tracker.sum(1)) # 0

You can test the efficiency of your solution with the following code. In this case too the code should finish almost immediately.

    tracker = DistanceTracker()
    total = 0
    for i in range(10**5):
        tracker.append(1)
        total += tracker.sum(1)
    print(total) # 4166749999583325000