Problem statement

https://leetcode.com/problems/count-substrings-with-only-one-distinct-letter/

Solution

What we can do is to apply groupby to calculate lengths of each group of adjacent equal elements. Then for each group we need to calculate i*(i+1)//2 and sum them all.

Complexity

Time complexity is O(n), space complexity as well. Space complexity can be reduced to O(1).

Code

class Solution:
    def countLetters(self, s):
        lens = [len(list(j)) for i, j in groupby(s)]
        return sum(i*(i+1)//2 for i in lens)