Queue reconstruction by height

Question

This is the "Queue Reconstruction by Height" problem from leetcode.com:

Suppose you have a random list of people standing in a queue. Each person is described by a pair of integers (h, k), where h is the height of the person and k is the number of people in front of this person who have a height greater than or equal to h. Write an algorithm to reconstruct the queue.

Note: The number of people is less than 1100.

Suppose I take only the tallest persons, all having the same maximum height. Their second values must be 0, 1, 2, 3... with no gaps at all, because they only count each other. Therefore, if there were no other persons at all, their second value must be their final index. What about the persons with second maximum height then? Suppose there are only tallest persons and just one more person who has slightly smaller height. What would be his position? Well, since he obviously only count tallest persons, his position would still be his second value. The next person of the same height counts only the previous person and all the tallest ones, but since they are all already in the queue, his second value would also be his index.

import collections

class Solution(object):
    def reconstructQueue(self, people):
        """
        :type people: List[List[int]]
        :rtype: List[List[int]]
        """
        d = collections.defaultdict(list)
        for height, count in people:
            d[height] += [count]
        
        recon_q = []
        for height, counts in sorted(d.items(), key = lambda x:x[0], reverse=True):
            for count in sorted(counts):
                recon_q.insert(count, [height, count])
        return recon_q

Is there any better data structure that can be used? I feel there must be a better way to do sorting here.

Note that the Solution class is required by LeetCode, even though it has no other function.

Answer 1

There are two difficulties here:

1. The input needs to be sorted in decreasing order by height, but increasing order by count. The code in the post handles this by collating the input by height and then sorting twice: first in reverse order by height, and then for a given height in normal order by count.

   This could be simplified by passing a key function to sorted so that it sorts in the required order. In particular, if take the key to be (-height, count) then the list will sort into the required order without any need for collation.

   def _queue_key(person):
       height, count = person
       return -height, count
   
   def reconstruct_queue_2(people):
       result = []
       for person in sorted(people, key=_queue_key):
           _, count = person
           result.insert(count, person)
       return result
   

2. Inserting an item into the middle a list takes time proportional to the length of the list (see the TimeComplexity page in the Python wiki). This means that the reconstruction algorithm takes time that's quadratic in the length of the queue. It would nice to do better than this.

   The data structure that's needed here is some kind of balanced tree. Python doesn't have anything suitable in the standard library, but the sortedcontainers package from the Python Package Index has a SortedListWithKey class that can do what we need.

   from sortedcontainers import SortedListWithKey
   
   def reconstruct_queue_3(people):
       result = SortedListWithKey(key=lambda _:0)
       for person in sorted(people, key=_queue_key):
           _, count = person
           result.insert(count, person)
       return result
   

   This is slightly tricksy code because we're not making use of the "sorted" property of the SortedListWithKey class — we use the key= argument to make all the items identical as far as sort order is concerned. We're using the class only to get efficient insertion at an arbitrary index in a list.