python: calculate column values ​based on previous rows and specified conditions

Question

I have table like below:

|input  | output|
|----   |  ---  |
|  5    |   0   |  
|  5    |   0   |
|  5    |   0   |
|  0    |   2   |
|  0    |   3   |
|  0    |   8   |
|  5    |   0   |
|  5    |   0   |

I need to add a "capacity" column, the values of which should be based on the conditions from other columns and the values of the previous "capacity" row.

Variable: START_CAPACITY = 20

Conditions for the first row:

if input > 0: START_CAPACITY + (input/4)
else: START_CAPACITY - (output/4)

Conditions for the rest rows:

if input > 0: capacity before + (input/4)
else: capacity before - (output/4)

Expected result:

| input | output |  capacity                       |
| ----- |  ----  |   ----                          |
|  5    |   0    | `START_CAPACITY` + (5/4) = 21.25|
|  5    |   0    | 21.25 + (5/4) = 22.50           |
|  5    |   0    | 22.50 + (5/4) = 23.75           |
|  0    |   2    | 23.75 - (2/4) = 23.25           |
|  0    |   3    | 23.25 - (3/4) = 22.50           |
|  0    |   8    | 22.50 - (8/4) = 20.50           |
|   5   |   0    | 20.50 + (5/4) = 21.75           |
|   5   |   0    | 21.75 + (5/4) = 23.00           |

Calculations in the capacity column are illustrative, for better understanding.

I'm a little fresh with python, tried the "cumsum" tricks, but in this case I don't know how to handle it.

Do you have any ideas for this? Thanks

Answer 1

I am not sure what you mean by "table", but when using lists, you could do something like this:

inputlist       = [5,5,5,0,0,0,5,5]
outputlist      = [0,0,0,2,3,8,0,0]
start_capacity  = 20

assert(len(inputlist) == len(outputlist))

capacity = []

for i in range(0, len(inputlist)):
    if i == 0:
        if inputlist[i] > 0:
            capacity.append(start_capacity + inputlist[i] / 4)
        else: 
            capacity.append(start_capacity - outputlist[i] / 4)
    else:
        if inputlist[i] > 0:
            capacity.append(capacity[i-1] + inputlist[i] / 4)
        else:
            capacity.append(capacity[i-1] - outputlist[i] / 4)
            
print('Capacity: ' + str(capacity))