Search code examples
pythonpandasdataframearray-broadcasting

Clever way to broadcast like a switch statement through a Pandas Dataframe


I have piecewise operation that I would do with a switch in other languages. Is there a way thru slicing or broadcasting that I can be efficient with doing something like the following? the data structure is a multi-index pandas Dataframe.

def spreadStat(data):
data.loc[data.Pval - data.Tval > 0 and data.Tval.value() > 75, 'Tval'] = data['Tval'] - 500
data.loc[data.Pval - data.Tval > 0 and data.Tval.value() < 25, 'Tval'] = data['Tval'] - 400
data.loc[data.Pval - data.Tval > 0 and 25 < data.Tval.value() < 40, 'Tval'] = data['Tval'] - 300
data.loc[data.Pval - data.Tval > 0 and 60 < data.Tval.value() < 75, 'Tval'] = data['Tval'] - 200
data.loc[data.Pval - data.Tval > 0 and 40 < data.Tval.value() < 60, 'Tval'] = data['Tval'] - 100

Seems like this could be a one-liner or some mapping of a function? Any help would be appreciated...


Solution

  • There is cut:

    bins = [-np.inf, 25, 40, 60, 75, np.inf]
    vals = [400, 300, 200, 100, 500]
    
    # I didn't test, but `astype` might not be necessary
    to_subtract = pd.cut(data.Tval, bins=bins, labels=vals).astype(float)
    
    data.loc[data.Pval > data.Tval, 'Tval'] -= to_subtract