Creating a pandas column of values with a calculation, but change the calculation every x times to a-CodePudding

I'm currently creating a new column in my pandas dataframe, which calculates a value based on a simple calculation using a value in another column, and a simple value subtracting from it. This is my current code, which almost gives me the output I desire (example shortened for reproduction):

subtraction_value = 3
data = pd.DataFrame({"test":[12, 4, 5, 4, 1, 3, 2, 5, 10, 9]} 

data['new_column'] = data['test'][::-1] - subtraction_value

When run, this gives me the current output:

print(data['new_column'])

[9,1,2,1,-2,0,-1,3,7,6]

However, if I wanted to use a different value to subtract on the column, from position [0], then use the original subtraction value on positions [1:3] of the column, before using the second value on position [4] again, and repeat this pattern, how would I do this iteratively? I realize I could use a for loop to achieve this, but for performance reasons I'd like to do this another way. My new output would ideally look like this:

subtraction_value_2 = 6

print(data['new_column'])

[6,1,2,1,-5,0,-1,3,4,6]

I hope this makes sense what I'm trying to achieve, and any help would be greatly appreciated.

CodePudding user response：

You can use positional indexing:

subtraction_value_2 = 6
col = data.columns.get_loc('new_column')
data.iloc[0::4, col] = data['test'].iloc[0::4].sub(subtraction_value_2)

or with numpy.where:

data['new_column'] = np.where(data.index%4,
                              data['test']-subtraction_value,
                              data['test']-subtraction_value_2)

output:

   test  new_column
0    12           6
1     4           1
2     5           2
3     4           1
4     1          -5
5     3           0
6     2          -1
7     5           2
8    10           4
9     9           6

CodePudding user response：

subtraction_value = 3
subtraction_value_2 = 6

data = pd.DataFrame({"test":[12, 4, 5, 4, 1, 3, 2, 5, 10, 9]})

data['new_column'] = data.test - subtraction_value
data['new_column'][::4] = data.test[::4] - subtraction_value_2

print(list(data.new_column))

Output:

[6, 1, 2, 1, -5, 0, -1, 2, 4, 6]