Np.Select add additional input type such as dict

Hi, I’m a huge fan of np.select() - it makes it very easy to categorize records in a table when there are a lot of potential categories. The only issue I have with it is that you have two separate lists to track, conditions and choices. This isnt so bad when dealing with only a handful of conditions but as the number grows it becomes difficult to track which list item in conditions corresponds to which list item in choices. I end up having to add a comment for each line of my conditions to say which choice it generates. I was wondering if we could alternately supply a dictionary/tuple (or any type that allows you to see the condition and choice in proximity to each other) where the keys are the choice and the values are the conditions.

Reproducing code example:

import numpy as np

#Current
conditions = [x == 1, x == 2, x == 3]
choices = ["One", "Two", "Three"] 
np.select(conditions, choices)

#Proposed
logic = {
    "One": x==1,
    "Two": x==2,
    "Three": x==3
}

np.select(logic)

[oliver-pola]

@matthewmturner If your problem is “end up having to add a comment for each line of my conditions to say which choice it generates”, you could use that dictionary for nice code in that part, and then call select like this?

np.select(logic.values(), logic.keys())

[matthewmturner]

@oliver-pola you’re totally right, thanks!