How can I refactor this code snippet to make it more efficient?

May 18, 2022

I’m trying to make a matrix of actions based off game theory. So if two people meet, they can either both share, or one steal, or both steal, etc.

The outline of what I have looks like this (not the whole matrix, just so you get an idea):

if first_agent.type == "AlwaysShare" and second_agent.type == "AlwaysShare":
    pass
elif first_agent.type == "AlwaysSteal" and second_agent.type == "AlwaysShare":
    pass
elif first_agent.type == "AlwaysShare" and second_agent.type == "AlwaysSteal":
    pass

Clearly this is very inefficient and prone to mistakes. How can I optimise this so I can efficiently manage interactions for the matrix?

>Solution :

When I have a situation with multiple options, I like to make a dictionary of functions that handle the different possibilities. The keys are the encoded and normalized inputs. In your case, you can generate a key like this:

def key(first_agent, second_agent):
    key = [first_agent.type, second_agent.type]
    key.sort()
    return tuple(key)

Then your dictionary would look something like this:

def handle_share():
    pass

def handle_one_steal():
    pass

def handle_both_steal():
    pass

# Etc.

action_map = {
    ('AlwaysShare', 'AlwaysShare'): handle_share,
    ('AlwaysShare', 'AlwaysSteal'): handle_one_steal,
    ('AlwaysSteal', 'AlwaysSteal'): handle_both_steal,
}

If you scope the functions carefully (e.g., make them methods or nested functions as necessary), you can fine-tune any side effects you may need them to have.

Now you can replace your if-block with something like:

action_map[key(first_agent, second_agent)]()

Use action_map.get instead if you have a suitable default or no-op in mind, and don’t want to get KeyError for potentially unknown interactions.

The nice thing about doing it this way is that you can easily add new interactions: just implement a function and register it in the dictionary. Want to have a third agent? Make key accept *args instead of a fixed number of agents. Need to define another type of agent? No problem. And so on.