Solution: Problem Challenge 2: Scheduling Tasks

from heapq import *
from collections import Counter

# We can calculate the solution using a two-part formula.
# Observe that we can calculate the length based on the maximum frequency of
# a task in the array. Then, we handle an edge case where we have multiple tasks
# with maximum frequency.

class Solution:
  def scheduleTasks(self, tasks, k):
    c = Counter(tasks)
    heap = [
      (-count, char) for char, count in c.items()
    ]
    heapify(heap)
    maxFreq = -heappop(heap)[0]
    maxFreqCount = 1
    while heap and -heap[0][0] == maxFreq:
      heappop(heap)
      maxFreqCount += 1
    
    if maxFreqCount > k:
      return len(tasks)
    
    # We will have a char, plus a gap of size k. So k+1.
    # We repeat this maxFreq - 1 times, as we don't need the last gap.