LC 0358 [H] Rearrange String k Distance Apart - ALawliet/algorithms GitHub Wiki

'''
Create a heapq orderd by frequency from big to small. The que length will be at most 26.
Create a cooldown dictionary, key being letter, value being remaining frequency, to hold letters in cool down state.

If k==0 or k==1, just return s. Otherwise,
Pop biggest frequency letter from que, add it to result.
Letter frequency -1, put it to the cool down dictionary.
For each iteration, pop letter that finished cooling down from the cool down dictionary, push it back to the frequency heap.
When it is impossible to make a valid result, the result array length must be shorter than input string. In such case we return "".
Time complexity O(nlgm), m is the number of distinct letter. Since k can only be 26, time complexity is O(n)
'''

class Solution:
    def rearrangeString(self, s: str, k: int) -> str:
        if k<=1: 
            return s
        d=collections.defaultdict(int)
        for c in s:
            d[c]+=1
        freqs=[[-d[k],k] for k in d]
        heapq.heapify(freqs)        
        cooling={}
        res=[]
        while freqs:
            freq,c=heapq.heappop(freqs)
            res.append(c)
            freq+=1
            if freq<0:                
                cooling[c]=[freq,c]
            if len(res)>=k and res[-k] in cooling:
                prevFreq,prevC=cooling.pop(res[-k])
                heapq.heappush(freqs,[prevFreq,prevC])        
        return ''.join(res) if len(res)==len(s) else ""