1 files changed, 10 insertions, 1 deletions
diff --git a/integer_sum_from_neighbours/generate-neighborsum.py b/integer_sum_from_neighbours/generate-neighborsum.py
index 7e06e02..14a8d23 100755
--- a/integer_sum_from_neighbours/generate-neighborsum.py
+++ b/integer_sum_from_neighbours/generate-neighborsum.py
@@ -1,9 +1,14 @@
 #!/usr/bin/env python
+# Usage for 4a (where 7 is the number of steps to run):
+#   $0 7 | z3 -smt -in
+# Usage for 4b:
+#   $0 7 b | z3 -smt -in
 import sys
 
 # i = 1..8
 n = 8
 steps = int(sys.argv[1])
+is_b = len(sys.argv) >= 3 and sys.argv[2] == 'b'
 
 literals = []
 preds = []
@@ -26,8 +31,12 @@ for i in range(1, n + 1):
     preds += ["(= a{i}_0 {i})".format(i=i)]
 
 # Goal: have a_3 = a_7 in any step:
+if not is_b:
+    goal = "(= a3_{i} a7_{i})"
+else:
+    goal = "(and (= a3_{i} a5_{i}) (= a5_{i} a7_{i}))"
 preds += [
-    "(or %s)" % " ".join("(= a3_{i} a7_{i})".format(i=i)
+    "(or %s)" % " ".join(goal.format(i=i)
     for i in range(steps + 1))
 ]