elegant code: itertools.groupby(S); 2. zip(count, count2)

809.expressive-words.py

# Time: O(n)
# Space: O(n)


# Before vs After 

#####################################
## Before
#####################################
 
'''
1. preprocesss S 

 23
aaabb
'''

class Solution(object):
    def expressiveWords(self, S, words):
        """
        :type S: str
        :type words: List[str]
        :rtype: int
        """
        # preprocess S 
        before_s, measure_s = self.slim(S)
        res = 0
        for word in words:
            before, measure = self.slim(word)
            if before == before_s:
                if self.compare(measure_s, measure):
                    res += 1
        return res 
            
    def slim(self, after):
        before = ""
        temp = after[0]
        measure = []
        for i in range(1, len(after)):
            if after[i] == after[i-1]:
                temp += after[i]
            else:
                before += temp[0]
                measure.append(len(temp))
                temp = after[i]
        
        before += temp[0]
        measure.append(len(temp))
            
        return before, measure
                  
    def compare(self, a, b):
        for i in range(len(a)):
            if a[i] < b[i]:
                return False 
            if a[i] > b[i] and a[i] == 2:
                return False 
        return True 

#####################################
## After
#####################################

class Solution(object):
    def expressiveWords(self, S, words):
        def slim(S):
            return zip(*[(k, len(list(grp)))
                         for k, grp in itertools.groupby(S)])

        R, count = slim(S)
        ans = 0
        for word in words:
            R2, count2 = RLE(word)
            if R2 != R: continue
            ans += all(c1 >= max(c2, 3) or c1 == c2
                       for c1, c2 in zip(count, count2))

        return ans