Update S4.py

diff --git a/Main/Python/2021/S4.py b/Main/Python/2021/S4.py
index 98cf4be61cb18c19353f1e2f44b0bc9f1b42b4a6..344b3925e4f9d51692360a486968149d6174c0cf 100644 (file)
--- a/Main/Python/2021/S4.py
+++ b/Main/Python/2021/S4.py
@@ -7,7 +7,7 @@
 #
 # Every day (D is total days), station X and Y swap
 #
-#
+# only want to get off the subway once
 #
 # """
 from queue import Queue
@@ -27,39 +27,32 @@ stations = list(map(int, input().split()))
 
 for i in range(D):
        X, Y = tuple(map(int, input().split()))
-       swap(X-1, Y-1)
-
-
+       swap(X - 1, Y - 1)
 
        q = Queue()
        q.put((0, 1))
 
-
-
        visited = set()
 
        while not q.empty():
                m, station = q.get()
 
-               if (m, station) in visited:
-                       continue
-               visited.add((m, station))
-
                if m > len(stations):
-                       continue
+                       break
 
                if station == N:
                        print(m)
                        break
-               # print(f"q:{q.qsize()}")
+
+
+               # take a train
+               if stations[m] == station:
+                       q.put((m + 1, stations[m + 1]))
 
                # take a walkway
+               if station in visited:
+                       continue
+               visited.add(station)
                if station in walkways:
                        for destination in walkways[station]:
                                q.put((m + 1, destination))
-
-               # take a train
-               if stations[m] == station and m + 1 < len(stations):
-                       q.put((m + 1, stations[m + 1]))
-               # wait at the station
-               q.put((m + 1, station))