How to replace a large if-elif statement by, for example, code using a dictionary?

Question

I have written the following code in Python, but the coding style looks bad and I would like to replace it with something nicer:

if A == 0 and B == 0:
    color = 'red'
elif A == 1 and B == 1:
    color = 'yellow'
elif A == 2 and B == 2:
    color = 'blue'            
elif A == 0 and B == 1:
    color = 'orange'
elif A == 1 and B == 0:
    color = 'orange'       
elif A == 2 and B == 1:
    color = 'green'
elif A == 1 and B == 2:
    color = 'green'
elif A == 0 and B == 2:
    color = 'purple'
elif A == 2 and B == 0:
    color = 'purple'  
        

I suggest using a dictionary as I wrote below would work, but I have difficulties in finding out how to code this in Python, because I have multiple values per key.

    color_dict = {
     "red": [[0,0]],
     "orange": [[0,1], [1,0]],
     "yellow": [[1,1]],
     "green": [[1,2],[2,1]],
     "blue": [[2,2]],        
     "purple": [[2,0],[0,2]]

Answer 1

Is there a reason you can't just use a 2d array to return the colors by indexing A and B ?

cols = [
    ['red',    'orange', 'purple'],
    ['orange', 'yellow', 'green' ],
    ['purple', 'green',  'blue'  ]
]

It could then be called as cols[A][B]

I just saw the other answer by eemz and realized the 2D array might get more complex and repetitive with more colors/options.

Answer 2

You could use a dict where the keys are tuples of (A, B)

color = {
    (0, 0): “red”,
    (1, 1): “yellow”,
    (2, 2): “blue”,
    etc etc
}
result = color[(a, b)]

Answer 3

Maybe you mean this one

It is not as efficient as the if-else block, but it could save you a lot of time typing. And also pay attention to the pattern matching with the in operator. It may not be as robust as you think.

def find_color(a, b):
    color_dict = {
        '0,0': 'red',
        '1,1': 'yellow',
        '2,2': 'blue',
        '0,1;1,0': 'orange',
        '2,1;1,2': 'green',
        '2,0;0,2': 'purple' 
    }
    maybe_key = list(filter(lambda key: f"{a},{b}" in key, color_dict.keys()))
    return color_dict[maybe_key[0]] if maybe_key else "NotFound" 

print(find_color(1, 2)) # green
print(find_color(2, 2)) # blue
print(find_color(3, 3)) # NotFound

Answer 4

Thanks for your feedback: I am using the 2D array solution now:

 A = 1
 B = 1
 cols = [
    ['red',    'orange', 'purple'],
    ['orange', 'yellow', 'green' ],
    ['purple', 'green',  'blue'  ]
 ] 
 color = colors_array[A][B]