How to call pandas columns with numeric suffixes in a for loop then apply conditions based on other-CodePudding

In python I am trying to update pandas dataframe column values based on the condition of another column value. Each of the column names have numeric suffixes that relate them. Here is a dataframe example:

Nodes_2 = pd.DataFrame([[0, 0, 37.76, 0, 0, 1, 28.32], [0, 0, 45.59, 0, 0, 1, 34.19], [22.68, 0, 22.68, 1, 0, 1, 34.02], [0, 0, 41.03, 0, 0, 1, 30.77], [20.25, 0, 20.25, 1, 0, 1, 30.37]], columns=['ait1', 'ait2', 'ait3', 'Type1', 'Type2', 'Type3', 'Flow'])

And the relevant 'Type' list:

TypeNums = [1, 2, 3, 4, 5, 6, 7, 8]

Specifically, I am trying to update values in the 'ait' columns with 'Flow' values if the 'Type' value equals 1; if the 'Type' value equals 0, the 'ait' value should be 0.

My attempt at applying these conditions is not working as it is getting hung up on how I am trying to reference the columns using the string formatting. See below:

for num in TypeNums:
if Nodes_2['Type{}'.format(num)] == 1:
    Nodes_2['ait{}'.format(num)] == Nodes_2['Flow']
elif Nodes_2['Type{}'.format(num)] == 0:
    Nodes_2['ait{}'.format(num)] == 0

That said, how should I call the columns with their numeric suffixes without typing duplicative code calling each name? And is this a correct way of applying the above mentioned conditions?

Thank you!

CodePudding user response：

The correct way is to use np.where or, in your case, just simple multilication:

for num in TypeNums:
    Nodes_2['ait{}'.format(num)] = Nodes_2['Type{}'.format(num)] * Nodes_2['Flow']

Or, you can multiply all the columns at once:

Nodes_2[['ait{}'.format(num) for num in TypeNums]] = Nodes_2[['Type{}'.format(num) for num in TypeNums]].mul(Nodes_2['Flow'], axis='rows').to_numpy()