Sort a list based on a given order

Viewed 11515

I have a list, say A = [[3,5],[1,3],[6,1]].

And I have another list, say B = [6,1,3]

I want to sort list A so that A becomes [[6,1],[1,3],[3,5]], which fits the given B — i.e. the first member of each A sublist should be sorted according to B.

5 Answers

derive a dict mapping the numbers in B to their indices and use that in the sort key function. This way you keep the key function constant time.

>>> A = [[3,5],[1,3],[6,1]]
>>> B = [6,1,3]
>>> srt = {b: i for i, b in enumerate(B)}
>>> sorted(A, key=lambda x: srt[x[0]])
[[6, 1], [1, 3], [3, 5]]

There are certainly a few caveats to this code. If a number were repeated in B you would get the index for the last entry used in the sort.

Also if there is an entry in A not matched in B, you have a KeyError. You could mitigate this somewhat by using dict.get with some default value, but if your input data was screwed to begin with, an error is a good thing to receive.

You can lookup where each [0] element of A is located using .index in B

>>> sorted(A, key = lambda i: B.index(i[0]))
[[6, 1], [1, 3], [3, 5]]

or in-place sort

>>> A.sort(key = lambda i: B.index(i[0]))
>>> A
[[6, 1], [1, 3], [3, 5]]

Assuming that A and B have the same length, each of the first elements of the tuples in A corresponds to an element in B and there are no duplicates, you can use this:

>>> A = [[3,5],[1,3],[6,1]]
>>> B = [6,1,3]
>>> A_ind = {b: [b, a] for (b, a) in A}
>>> [A_ind[b] for b in B]
[[6, 1], [1, 3], [3, 5]]

This simply creates a dict, mapping the b part from A to the whole elements and looking up the respective elements in B in that dict. Using neither sorted not index, this has complexity of O(n).

If there can be duplicate elements, you can create a dict (or defaultdict) mapping keys to lists of elements having that key:

>>> A = [[3,5],[1,3],[6,1],[1,4]]
>>> A_ind = collections.defaultdict(list)
>>> for b, a in A:
...     A_ind[b].append([b,a])
...     
>>> B = [6,1,3,1]
>>> [a for b in B for a in A_ind[b]]
[[6, 1], [1, 3], [1, 4], [3, 5], [1, 3], [1, 4]]

Or turn the list into iterators to distribute the values in A to the matching keys in B:

>>> A_ind = {k: iter(v) for k, v in A_ind.items()}
>>> [next(A_ind[b]) for b in B]
[[6, 1], [1, 3], [3, 5], [1, 4]]

This is very similar to CoryKramer's answer but it is a bit faster. Instead of using the sort function, manually insert items to a another list. This way the index method gets called less often. (This code assumes the elements in b are distinct).

a= [[3,5],[1,3],[6,1]]
b=[6,1,3]

c= [-1]*len(b) # dummy data

for x in a:
    c[b.index(x[0])]= x

print(c)
sorted(A, key=lambda x: B.index(x[0]))
Related