LeetCode 1086: High Five

link

We create a map: {id -> minq(score)} and we cap the length of the minq at 5. In that way, in the end, we have the top five scores per id.

Say len(items) = n. There are two extreme cases:

  1. All scores belong to a single id
    • We would have n heap operations. But the heap would have fixed length, so time per heap operation is \mathcal{O}(1).
    • The dict has a single item, so sorting takes time \mathcal{O}(1).
    • Total space is \mathcal{O}(1).
  2. \frac{n}{5} unique ids
    • We would have \frac{n}{5} keys in the dict. No heap operation is involved. The time complexity per dict item is still \mathcal{O}(1).
    • Sorting the dict takes time \mathcal{O}(n \cdot \lg{n}).
    • Space is \mathcal{O}(n).

Time: \mathcal{O}(n \cdot \lg{n}), space: \mathcal{O}(n).

from heapq import heappush as push, heappop as pop

class Solution:
    def highFive(self, items: List[List[int]]) -> List[List[int]]:
        hi_five_map = {}
        for id, score in items:
            if id not in hi_five_map:
                hi_five_map[id] = [score]
                continue

            push(hi_five_map[id], score)
            if len(hi_five_map[id]) > 5:
                pop(hi_five_map[id])
        
        hi_fives = []
        for id in sorted(hi_five_map):
            top_scores = hi_five_map[id]
            hi_fives.append( [id, sum(top_scores) // len(top_scores)] )

        return hi_fives

Note, the heap plays a role only if there are more than five scores for an id. So, we can optimize by keeping a list of scores if the count of scores is 5 or less and switch to heap once the count of scores go beyond 5.

from heapq import heapify, heappush as push, heappop as pop

class ScoreHeap:
    def __init__(self, score):
        self.is_heap = False
        self.q = [score]
    
    def push(self, score):
        if self.is_heap:
            if score <= self.q[0]:
                return
            push(self.q, score)
            pop(self.q)
            return

        self.q.append(score)
        if len(self.q) > 5:
            heapify(self.q)
            pop(self.q)
            self.is_heap = True

    def average(self):
        return sum(self.q) // len(self.q)

class Solution:
    def highFive(self, items: List[List[int]]) -> List[List[int]]:
        hi_five_map = {}
        for id, score in items:
            if id not in hi_five_map:
                hi_five_map[id] = ScoreHeap(score)
                continue

            hi_five_map[id].push(score)
        
        hi_fives = []
        for id in sorted(hi_five_map):
            hi_fives.append([id, hi_five_map[id].average()])

        return hi_fives

Leave a comment