From 2251f2e6f75c6a0078b4cb2b0da0b3f00e3f063a Mon Sep 17 00:00:00 2001 From: Emmy D'Anello Date: Mon, 13 Feb 2023 00:24:43 +0100 Subject: [PATCH] Calculate routes with IATA codes --- app.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/app.py b/app.py index 21079e5..be10657 100644 --- a/app.py +++ b/app.py @@ -166,19 +166,19 @@ def find_routes(day, origin, destination): valid_routes = [] for train in tqdm(trains): - if train.orig == origin: + if train.orig_iata == origin: it = [train] - if train.dest == destination: + if train.dest_iata == destination: # We hope that we have a direct train valid_routes.append(it) else: explore.append(it) - per_arr_explore.setdefault(train.dest, []) - per_arr_explore[train.dest].append(it) + per_arr_explore.setdefault(train.dest_iata, []) + per_arr_explore[train.dest_iata].append(it) continue - for it in list(per_arr_explore.get(train.orig, [])): - if any(train.dest == tr.dest or train.dest == origin for tr in it): + for it in list(per_arr_explore.get(train.orig_iata, [])): + if any(train.dest_iata == tr.dest_iata or train.dest_iata == origin for tr in it): # Avoid loops continue @@ -186,13 +186,13 @@ def find_routes(day, origin, destination): if last_train.arr <= train.dep: new_it = it + [train] - if train.dest == destination: + if train.dest_iata == destination: # Goal is achieved valid_routes.append(new_it) else: explore.append(new_it) - per_arr_explore.setdefault(train.dest, []) - per_arr_explore[train.dest].append(new_it) + per_arr_explore.setdefault(train.dest_iata, []) + per_arr_explore[train.dest_iata].append(new_it) return valid_routes