FoodDelivery: improve solution for b

Now requires truck to return to a location, also accept any location
instead of just the supply location S.

1 files changed, 8 insertions, 3 deletions

diff --git a/food_delivery/food-to-table.py b/food_delivery/food-to-table.py
index b94ad0c..c66a774 100755
--- a/food_delivery/food-to-table.py
+++ b/food_delivery/food-to-table.py
@@ -7,6 +7,8 @@ n = 5
 
 matrix = []
 nextNodes = {}
+# The pair that shows a loop (for b)
+pair_for_b = [None] * 2
 
 for line in sys.stdin:
     line = line.strip()
@@ -23,6 +25,10 @@ for line in sys.stdin:
             matrix.append([None] * n)
 
         matrix[step][villagenr] = value
+    elif line.startswith('k'):
+        m = re.match(r'k(\d+) -> (\d+)', line)
+        step, value = map(int, m.groups())
+        pair_for_b[step] = value
     else:
         m = re.match(r'l(\d+) -> (\d+)', line)
         step, value = map(int, m.groups())
@@ -45,8 +51,7 @@ for step, cells in enumerate(matrix):
         else:
             line += " & %5d" % value
     line += r" & %s \\" % "SABCD"[nextNodes[step]]
-    if step > 0:
-        if all(x >= y for x, y in zip(matrix[step][1:], matrix[0][1:])):
-            line += " % HIGHER"
+    if any(x == step for x in pair_for_b):
+        line += " % PAIR"
     #line += r" \\"
     print(line)