Solve puzzle 21 part 2

1 files changed, 179 insertions, 13 deletions

diff --git a/21.c b/21.c
index b1eebd4..34c648f 100644
--- a/21.c
+++ b/21.c
@@ -1,11 +1,13 @@
 
 //#define _GNU_SOURCE
 
+#include <limits.h>
 #include <regex.h>
 #include <search.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <unistd.h>
 
 
 #define arrlen(x) (sizeof(x)/sizeof((x)[0]))
@@ -17,6 +19,8 @@ enum operation {
 	OP_SUB,
 	OP_MUL,
 	OP_DIV,
+	OP_CMP,
+	OP_UNK,
 };
 
 
@@ -28,11 +32,15 @@ operation_get(char op)
 	case '-': return OP_SUB;
 	case '*': return OP_MUL;
 	case '/': return OP_DIV;
+	case '=': return OP_CMP;
+	case '?': return OP_UNK;
 	}
 	return OP_VAL;
 }
 
 
+#define MONKEY_UNKNOWN LONG_MIN
+
 struct monkey {
 	char name[5];
 	enum operation op;
@@ -118,36 +126,44 @@ monkey_find(void * monkey_tree, char const * name)
 long
 monkey_value(void * monkey_tree, struct monkey * m)
 {
-	long v;
+	long left, right, v;
 
 	if (!m) {
 		fprintf(stderr, "%s: Bad monkey\n", __FUNCTION__);
 		return 0;
 	}
 
+	if (m->op == OP_VAL)
+		return m->value;
+
+	if (m->op == OP_UNK)
+		return MONKEY_UNKNOWN;
+
+	left = monkey_value(monkey_tree, monkey_find(monkey_tree, m->left));
+	right = monkey_value(monkey_tree, monkey_find(monkey_tree, m->right));
+	if (left == MONKEY_UNKNOWN || right == MONKEY_UNKNOWN)
+		return MONKEY_UNKNOWN;
+
 	switch (m->op) {
-	case OP_VAL:
-		v = m->value;
+	case OP_CMP:
+		// If they match, return that value, otherwise MONKEY_UNKNOWN
+		v = left == right ? left : MONKEY_UNKNOWN;
 		break;
 
 	case OP_ADD:
-		v = monkey_value(monkey_tree, monkey_find(monkey_tree, m->left))
-			+ monkey_value(monkey_tree, monkey_find(monkey_tree, m->right));
+		v = left + right;
 		break;
 
 	case OP_SUB:
-		v = monkey_value(monkey_tree, monkey_find(monkey_tree, m->left))
-			- monkey_value(monkey_tree, monkey_find(monkey_tree, m->right));
+		v = left - right;
 		break;
 
 	case OP_MUL:
-		v = monkey_value(monkey_tree, monkey_find(monkey_tree, m->left))
-			* monkey_value(monkey_tree, monkey_find(monkey_tree, m->right));
+		v = left * right;
 		break;
 
 	case OP_DIV:
-		v = monkey_value(monkey_tree, monkey_find(monkey_tree, m->left))
-			/ monkey_value(monkey_tree, monkey_find(monkey_tree, m->right));
+		v = left / right;
 		break;
 
 	default:
@@ -159,13 +175,129 @@ monkey_value(void * monkey_tree, struct monkey * m)
 }
 
 
+// Find a monkey with OP_UNK in the tree, fixup its value to match the value passed in, and return that monkey.
+struct monkey *
+monkey_fixunknown(void * monkey_tree, struct monkey * m, long value)
+{
+	struct monkey * mleft, * mright, * munk;
+	long vleft, vright, v;
+
+	if (!m) {
+		fprintf(stderr, "%s: Bad monkey\n", __FUNCTION__);
+		return NULL;
+	}
+
+	if (m->op == OP_UNK) {
+		// Is an unknown value. Set it to the value we passed in and return it!
+		m->op = OP_VAL;
+		m->value = value;
+		return m;
+	}
+
+	if (m->op == OP_VAL)
+		// Can't fix an existing value.
+		return NULL;
+
+
+	// Binary operation. Get the left and right monkeys/values
+	mleft = monkey_find(monkey_tree, m->left);
+	mright = monkey_find(monkey_tree, m->right);
+
+	vleft = monkey_value(monkey_tree, mleft);
+	vright = monkey_value(monkey_tree, mright);
+
+	if (vleft == MONKEY_UNKNOWN && vright == MONKEY_UNKNOWN)
+		// Two unknowns. Can't fix.
+		return NULL;
+
+	if (vleft != MONKEY_UNKNOWN && vright != MONKEY_UNKNOWN)
+		// No unknowns. Nothing to fix.
+		return NULL;
+
+	// Work out which is the unknown, and which is the known
+	if (vleft == MONKEY_UNKNOWN) {
+		munk = mleft;
+		v = vright;
+	}
+	else {
+		munk = mright;
+		v = vleft;
+	}
+
+	if (m->op == OP_CMP) {
+		// It's a comparison. If the passed in value is different from
+		// the known value, we can't fix.
+		if (value != MONKEY_UNKNOWN && value != v)
+			return NULL;
+		// OK, fix the unknown side.
+		return monkey_fixunknown(monkey_tree, munk, v);
+	}
+
+	if (value == MONKEY_UNKNOWN)
+		// Trying to fix a binary operation, where the result we need
+		// to produce is unknown?
+		return NULL;
+
+	// OK, we have a binary op where we know the result (#value) and one of
+	// the operands, but not the other. Recursively fix the operand we don't
+	// know.
+	switch (m->op) {
+	case OP_ADD:
+		// unknown + known = value => unknown = value - known
+		return monkey_fixunknown(monkey_tree, munk, value - v);
+
+	case OP_SUB:
+		if (vleft == MONKEY_UNKNOWN)
+			// unknown - right = value => unknown = value + right
+			return monkey_fixunknown(monkey_tree, munk, value + vright);
+		else
+			// left - unknown = value => unknown = left - value
+			return monkey_fixunknown(monkey_tree, munk, vleft - value);
+
+	case OP_MUL:
+		// unknown * known = value => unknown = value / known
+		return monkey_fixunknown(monkey_tree, munk, value / v);
+
+	case OP_DIV:
+		if (vleft == MONKEY_UNKNOWN)
+			// unknown / right = value => unknown = value * right
+			return monkey_fixunknown(monkey_tree, munk, value * vright);
+		else
+			// left / unknown = value => unknown = left - value
+			return monkey_fixunknown(monkey_tree, munk, vleft / value);
+
+	default:
+		fprintf(stderr, "Unexpected operation for %s: %d\n", m->name, m->op);
+		break;
+	}
+
+	return NULL;
+}
+
+
 int
-main()
+main(int argc, char ** argv)
 {
+	int part = 1, i;
 	regex_t reval, reop;
 	char buf[BUFSIZ];
 	void * monkey_tree = NULL;
+	struct monkey * m;
+
+
+	// Parse command options
+	while ((i = getopt(argc, argv, "p:")) != -1) {
+		switch (i) {
+		case 'p':
+			part = atoi(optarg);
+			break;
 
+		case -1:
+			return -1;
+		}
+	}
+
+	// Compile regexes
 	if (regcomp(&reval, "([[:alpha:]]{4}): (-?[[:digit:]]+)", REG_EXTENDED) != 0
 			|| regcomp(&reop, "([[:alpha:]]{4}): ([[:alpha:]]{4}) ([-+*/]) ([[:alpha:]]{4})", REG_EXTENDED) != 0)
 	{
@@ -173,6 +305,7 @@ main()
 		return -1;
 	}
 
+	// Read file.
 	while (fgets(buf, sizeof(buf), stdin)) {
 		regmatch_t matches[5];
 		struct monkey * m;
@@ -192,18 +325,51 @@ main()
 		}
 		else {
 			fprintf(stderr, "Bad monkey: %s\n", buf);
+#ifdef _GNU_SOURCE
+			tdestroy(monkey_tree, free);
+#endif
 			regfree(&reop);
 			regfree(&reval);
 			return -1;
 		}
 
+		if (part == 2 && strcmp(m->name, "root") == 0) {
+			m->op = OP_CMP;
+			m->value = MONKEY_UNKNOWN;
+		}
+		else if (part == 2 && strcmp(m->name, "humn") == 0) {
+			m->op = OP_UNK;
+			m->value = MONKEY_UNKNOWN;
+		}
+
 		tsearch(m, &monkey_tree, monkey_cmp_t);
 	}
 
+	// Output result.
+	m = monkey_find(monkey_tree, "root");
+	switch (part) {
+	case 1:
+		printf("root value: %ld\n", monkey_value(monkey_tree, m));
+		break;
 
+	case 2:
+		if ((m = monkey_fixunknown(monkey_tree, m, MONKEY_UNKNOWN)) == NULL)
+			printf("Unknown value not found\n");
+		else
+			printf("Fixed unknown value for %s: %ld\n", m->name, m->value);
+		break;
 
-	printf("root value: %ld\n", monkey_value(monkey_tree, monkey_find(monkey_tree, "root")));
+	default:
+		fprintf(stderr, "Unexpected part %d\n", part);
+#ifdef _GNU_SOURCE
+		tdestroy(monkey_tree, free);
+#endif
+		regfree(&reop);
+		regfree(&reval);
+		return -1;
+	}
 
+	// Tidy and done.
 #ifdef _GNU_SOURCE
 	tdestroy(monkey_tree, free);
 #endif