Same character

A substring is a contiguous string inside a string. For example, the substrings of the string abc are a, b, c, ab, bc and abc. Your task is to count how many substrings have the same character at all positions.

The time complexity of the algorithm should be O(n).

In a file samechar.py, implement a function count that returns the desired count.

def count(s):
    # TODO

if __name__ == "__main__":
    print(count("aaa")) # 6
    print(count("abbbcaa")) # 11
    print(count("abcde")) # 5