Optimise a function with numerous conditions that depends on the previous row in a Python dataframe

per pandas iteration guidance

You should never modify something you are iterating over. This is not guaranteed to work in all cases. Depending on the data types, the iterator returns a copy and not a view, and writing to it will have no effect!

suggested

from typing import Iterable, Tuple import pandas as pd COLS = ['country_ID', 'ID', 'direction', 'date'] DATA = [['ESP_1', 0, 'IN', '2021-02-28'], ['ENG', 0, 'IN', '2021-03-03'], ['ENG', 0, 'OUT', '2021-03-04'], ['ESP_2', 0, 'IN', '2021-03-05'], ['FRA', 1, 'OUT', '2021-03-07'], ['ENG', 1, 'OUT', '2021-03-09'], ['ENG', 1, 'OUT', '2021-03-10'], ['ENG', 2, 'IN', '2021-03-13']] def both_in(country_id: str): """where both condtions were `IN`""" esp, eng, fra = (0, 0, 0) if country_id == 'FRA': error_code = 0 elif country_id in ('ESP_1', 'ESP_2'): error_code = 2 fra = 1 else: error_code = 3 return (esp, eng, fra, f'ERROR:{error_code}') def both_out(country_id: str): """where both contionds were `OUT`""" esp, eng, fra = (0, 1, 0) if country_id == 'ENG': error_code = 8 elif country_id in ('FRA', 'ESP_2'): error_code = 9 fra = 1 else: error_code = 10 esp, fra = 1, 1 return (esp, eng, fra, f'ERROR:{error_code}') def in_out(country_id: str): """where current was IN and previous was OUT""" esp, eng, fra = (0, 0, 0) if country_id == 'ENG': error_code = 0 elif country_id in ('FRA', 'ESP_2'): error_code = 4 fra = 1 else: error_code = 5 esp = 1 fra = 1 return (esp, eng, fra, f'ERROR:{error_code}') def out_in(country_id: str): """where current was `OUT` and previous was `IN`""" esp, eng, fra = (0, 0, 0) if country_id == 'ENG': error_code = 0 elif country_id in ('ESP_1', 'ESP_2'): error_code = 6 eng, fra = 1, 1 else: error_code = 7 eng = 1 return (esp, eng, fra, f'ERROR:{error_code}') def create_columns_analysis(df: pd.DataFrame)->Iterable[Tuple[int,int,int,str,pd.Timestamp]]: """create_columns_analysis""" # in your logic there are 4 potential driving conditions based on the direction # of the current row and the previous row. so we'll make a dictionary that we can # index and call the associated functions. direction = { ('IN', 'IN'): both_in, ('OUT', 'OUT'): both_out, ('IN', 'OUT'): in_out, ('OUT', 'IN'): out_in, } # to align the direction slice # - the last row # - the first row # - and zip them together def iter_countires(): """yields a (ESP,ENG,FRA,ERROR,MEAN_TIME)""" list_ids = [] # because we sliced the first row from the loop yield that value first def first_row(country_id): if country_id == 'ENG': return (0, 1, 0, 'ERROR:1', pd.NA) return (0, 0, 0, 'ERROR:1', pd.NA) yield first_row(df['country_ID'][0]) for previous, current in zip(df[:-1].itertuples(), df[1:].itertuples()): time_delta = current.date-previous.date # If it is the first time the ID is detected if current.country_ID == 'ENG' and current.ID not in list_ids: list_ids.append(current.ID) yield (0, 1, 0, 'ERROR:1', time_delta) elif current.country_ID == 'ENG': # indexing our dict with the ('IN','OUT') to get the function conditional_func = ( direction[(previous.direction, current.direction)]) # call the function and pass the previous.country_ID as thats the only var it relies on # unpack thoes values and tack on the timedelta yield (*conditional_func(previous.country_ID), time_delta) else: # you could create a different conditional func dict if you wanted to # handle country_ID logic differently yield (0, 0, 0, pd.NA, pd.NA) return iter_countires() def start(): """start""" df = pd.DataFrame(DATA, columns=COLS) df['date'] = pd.to_datetime(df['date']) df[['visit_ESP', 'visit_ENG', 'visit_FRA', 'error', 'mean_time']]=( tuple(create_columns_analysis(df))) print(df) if __name__ == '__main__': start() 

result

 country_ID ID direction date visit_ESP visit_ENG visit_FRA error mean_time 0 ESP_1 0 IN 2021-02-28 0 0 0 ERROR:1 <NA> 1 ENG 0 IN 2021-03-03 0 1 0 ERROR:1 3 days 00:00:00 2 ENG 0 OUT 2021-03-04 0 0 0 ERROR:0 1 days 00:00:00 3 ESP_2 0 IN 2021-03-05 0 0 0 <NA> <NA> 4 FRA 1 OUT 2021-03-07 0 0 0 <NA> <NA> 5 ENG 1 OUT 2021-03-09 0 1 0 ERROR:1 2 days 00:00:00 6 ENG 1 OUT 2021-03-10 0 1 0 ERROR:8 1 days 00:00:00 7 ENG 2 IN 2021-03-13 0 1 0 ERROR:1 3 days 00:00:00