Python how to sort list of objects based on their attribute according to another list

Question

That is a confusing title.

Lets say I have a list of IDs

ids = [3, 1, 2, 4]

And I am retrieving the objects in one go, for example in Django:

records = Record.objects.filter(id__in=ids)

How do I keep the order of records according to another list, elegantly , preferably using built-in sort functions ?

The only way I can think of right now is in traditional code with minimum loops, covert the records to dictionary first, then loop through the ids to pick the record in the dictionary directly, see below:

result = []
dict_of_records = {r.id: r for r in records}
for id in ids:
    result.append(dict_of_records[id])

Answer 1

Use sorted() and specify the sort key with the key keyword arg:

sorted(records, key=lambda x: ids.index(x))

this however will be inefficient , because it will require a lookup ( ids.index(x) ) for each record.

It will be more efficient if you precompute the sort key:

sortkey = {j:i for i,j in enumerate(ids)}
sorted(records, key=lambda x: sortkey[x])

Answer 2

Personally, I wouldn't use a sort-based solution. A sort would take O(Nlog(N)) time for N records, but we can solve the problem in O(N) time. That said, a sort would look like this:

positions = {j:i for i, j in enumerate(ids)}
result = sorted(Record.objects.filter(id__in=ids), key=lambda r:positions[r.id])

What I would do is use the positions table to put each record in its final position immediately:

result = [None]*len(positions)
for record in Record.objects.filter(id__in=ids):
    result[positions[record.id]] = record