More elegant list comprehension using another enumerated list as a key

Viewed 40

I have a list of lists:

>>> array = [["A","B","C"],["C","B","A"]]

I also have a key which I used enumerate to create:

>>> key = list(enumerate(["A","B","C"]))
>>> print (key)
[(0,"A"),(1,"B"),(2,"C")]

I want to use list comprehension to convert the strings in the array into numbers based on the key. I did the following and was able to get the result I wanted:

>>> for i in key:
...    board = [[i[0] if x in i else x for x in y] for y in board] 
>>> print(board)
[[0,1,2],[2,1,0]]

However, I can't seem to figure out how to use a single list comprehension without a for loop to obtain the result. Any ideas on how I would be able to do this?

4 Answers

If I understand you correctly, you want to convert the for-loop into list comprehension. You can create a mapped (in your question this is inverted key dictionary`) and then use list comprehension:

array = [["A","B","C"],["C","B","A"]]
mapper = {v:k for k, v in dict(enumerate(["A","B","C"])).items()}

print( [[mapper[v] for v in subl] for subl in array] )

Prints:

[[0, 1, 2], [2, 1, 0]]

EDIT: Thank to @kaya3, lighter version with just enumerate:

array = [["A","B","C"],["C","B","A"]]
mapper = {v:k for k, v in enumerate(["A","B","C"])}

print( [[mapper[v] for v in subl] for subl in array] )

You can use the defaultdict length trick to get unique integer representations by order of occurrence in a single iteration:

from collections import defaultdict

d = defaultdict(lambda: len(d))

array = [["A", "B", "C"], ["C", "B", "A"]]
array = [[d[x] for x in sub] for sub in array]
# [[0, 1, 2], [2, 1, 0]]

With python 3.8 you can use the same approach as @AndrejKesely in one line, creating the mapper via walrus-operator inside the list comprehension:

array = [["A","B","C"],["C","B","A"]]

arr2 = [[ mapper[p] for p in l] for l in array 
        if (mapper := {v:k for k,v in enumerate(array[0])})] 

print(arr2)

Output:

[[0, 1, 2], [2, 1, 0]]

You can make key into a dictionary mapping letters to their indices, then do the mapping:

letters = ['A', 'B', 'C']
array = [['A', 'B', 'C'], ['C', 'B', 'A']]
key = {x: i for i, x in enumerate(letters)}
array = [list(map(key.get, row)) for row in array]

Alternatively, since letters is quite short, letters.index is efficient enough that there is not much point building a key at all:

letters = ['A', 'B', 'C']
array = [['A', 'B', 'C'], ['C', 'B', 'A']]
array = [list(map(letters.index, row)) for row in array]
Related