Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python - too many 'elif: return()' statements?

I have ths function that's passed two arguments:

def function(a, b):
    if   a == 'd' : return(4*b)
    elif a == '!' : return(5)
    elif a == 'k' : return(1-2*k+k**2)
    elif a == 'Z' : return(1/k)
    (...)

a is checking it's equality to a single character, and b is always a number; the function always returns a number as well. Sometimes it's not always a simple return though.

def function(a, b):
    (...)
    elif a == '2':
        temp_b = foo(b)
        if b == 2 : temp_b += 2
        return(temp_b)

I have a very long list of elif statements, is there a better way to do this?

like image 884
Graviton Avatar asked Jun 29 '17 03:06

Graviton


2 Answers

Actually yes. First off, Python doesn't have the iconic switch statement, which provides a binary search when the statement becomes long enough, which is much faster than a linear search.

In Python, you can use a dictionary, either with exact values, or custom functions:

def other_case(x):
    '''We can store non-lambdas too'''

    return 8

functions = {
    'd': lambda x: 4*x,
    '!': lambda x: 5,
    'k': lambda x: (1-2*x+x**2),
    'Z': lambda x: (1/x),
    '*': other_case,
}

To call this, just right a short wrapper:

def call(a, x):
    return functions[a](x)

This has O(1) time, or constant complexity, so in addition to being a lot more readable, it will be a lot faster too.

Edit

If you have numbers of a specific range, that can be all over that entire range, you can also use a list and a transformation. Say I want to process something from a mass of 2000 Da to 5000 Da (I'm a biologist), at intervals of 100 Da. It makes no sense to encode a list of 500 items, but I could use a 30-item list to consider the entire range.

from __future__ import division         # for //, floor division

def mass2000(x):
    '''Do something for mass of 2000'''

    return 1/x

def mass2100(x):
    '''Do something for mass of 2100'''

    return x

def mass2200(x):
    '''Do something for mass of 2200'''

    return x**2


lookup = [
    mass2000,
    mass2100,
    mass2200,
    # ....
]


def call(mass, x):
    if (mass < 2000 or mass > 5000):
        raise ValueError("Mass out of range")

    return lookup[(mass - 2000) // 100](x)
like image 81
Alexander Huszagh Avatar answered Nov 15 '22 04:11

Alexander Huszagh


It could be beneficial to separate all of your cases into individual mini functions or lambda functions, then have a dictionary with the key being the character with the correlating function/lambda function as the value.

Here's an example:

def function_2(b):
    temp_b = foo(b)
    if b == 2 : temp_b += 2
    return(temp_b)

fp = {
    2 : function_2
}

Then you can have a loop that loops through the dictionary and matches the keys and once matched you can pass b in like this.

for element in fp:
    if element == a:
        fp[element](b)
like image 22
glyif Avatar answered Nov 15 '22 03:11

glyif