Optimize if tree in Python

Question

I have an "if decision tree" and I want to know if it's possible to optimize it:

    def s(a, b):
    """
    :param a: 0,1,2
    :param b: 0,1,2
    :return: 0,1,2
    """
    if a == 0:
        if b == 0:
            return 2
        elif b == 1:
            return 2
        else:  # b == 2
            return 0
    elif a == 1:
        if b == 0:
            return 2
        elif b == 1:
            return 2
        else:  # b == 2
            return 1
    else:  # a==2
        if b == 0:
            return 0
        elif b == 1:
            return 1
        else:  # b==2
            return 2

All cases included, I'm trying use (a,b) == (1,2) or (a,b) == (2,1) return 1 and the same when return 0. Other cases return 2, but it is slower.

EDIT: I have tested (time and valid (all valid)) all proposed "ifs" and: s is my explicit function ss is first proposed function sss is second proposed function etc. and here I have profile:

You could optimize using and and or but if you're concerned about optimizing an if-statement, you're usually over thinking optimization — OneCricketeer
– OneCricketeer, Commented Oct 24, 2015 at 15:02
@cricket_007 I disagree. If the logic is really as simple as OP wrote in the last sentence, you should be able to represent exactly that logic. Writing out all possible cases is not a good way to get the logic across. And it’s very error prone, since I would have to look many times to verify that the cases are really all correct. — poke
– poke, Commented Oct 24, 2015 at 15:04
I wrote it this way only because I need to check if all possibilities return valid value. Then decide to ask stack community :). — Behoston
– Behoston, Commented Oct 25, 2015 at 6:26

poke · Accepted Answer · 2015-10-24 15:08:57Z

5

Just check those values explicitely by comparing the tuples, and return 2 in all other cases:

v = (a, b)
if v == (1, 2) or v == (2, 1):
    return 1
elif v == (0, 2) or v == (2, 0):
    return 0
else:
    return 2

You could even go a bit further and think about what your check does: If any of the numbers is two, you return the other number.

if a == 2:
    return b
elif b == 2:
    return a
else:
    return 2

In general, of course you could also create an exact mapping, e.g. using a dictionary, that allows you to look up the result directly (instead of having a complex if/else structure). But when your logic to determine the result really fits in a single sentence, then you should really try to implement it in a way that keeps that concise logic intact. Otherwise, it would be very difficult to spot errors (for example, it took me a while to verify that your if/else code really does what you said it would). Of course, you can always use those different possible combinations in a unit test to make sure that your logic actually covers it all.

edited Oct 24, 2015 at 15:08

answered Oct 24, 2015 at 15:01

poke

392k80 gold badges596 silver badges632 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Behoston Over a year ago

your second function is a little faster than explicit all if statement

Randy · Accepted Answer · 2015-10-24 15:02:22Z

2

If you know you're not going to get a=3 or b=-1, here's a oneliner:

def new_s(a,b):
    return min(a,b) if max(a,b) == 2 else 2

answered Oct 24, 2015 at 15:02

Randy

14.9k3 gold badges40 silver badges44 bronze badges

2 Comments

Behoston Over a year ago

this is the slowest function

Randy Over a year ago

If you're talking about optimizing from a speed perspective, then you're better off doing something with numpy, Cython, or similar tools rather than pure Python.

blackmamba · Accepted Answer · 2015-10-25 07:43:14Z

0

I would prefer to use a dictionary.

mydict[0][0] = 2
mydict[0][1] = 2
mydict[0][2] = 0
mydict[1][0] = 2
mydict[1][1] = 2
mydict[1][2] = 1
mydict[2][0] = 0
mydict[2][1] = 1
mydict[2][2] = 2

print mydict[a][b]

How to create dict of dict:

mydict = dict()
mydict[0] = dict()
mydict[0][1] = 2

edited Oct 25, 2015 at 7:43

answered Oct 24, 2015 at 15:02

blackmamba

2,00211 gold badges37 silver badges61 bronze badges

3 Comments

Behoston Over a year ago

How to make two integer key dictionary? It's dict of dicts? I use tuple keys but I'm not sure if that is appropriate

blackmamba Over a year ago

Yes, that is dict of dict. Added a sample to the code.

poke Over a year ago

@Behoston You can use tuple keys too, that’s fine.

Collectives™ on Stack Overflow

Optimize if tree in Python

3 Answers 3

1 Comment

2 Comments

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

2 Comments

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related