One character

Your task is to count the number of substrings in a string where all the characters are the same. For example, the string abbbcaa has 11 such substrings:

• a (three times)
• aa
• b (three times)
• bb (twice)
• bbb
• c

You may assume that the string consists of the characters a–z and contains at most 10^5 characters. The goal is an algorithm with time complexity O(n).

In a file onechar.py, implement a function count that returns the number of substrings.

def count(s):
    # TODO

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