path: root/16.c

#include <regex.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>


#define arrlen(x) (sizeof(x)/sizeof((x)[0]))


void *
qelem(void * base, size_t size, int i)
{
	return (char *) base + i * size;
}


void
qswap(void * base, size_t size, int i, int j)
{
	char tmp[size];
	memcpy(tmp, qelem(base, size, i), size);
	memcpy(qelem(base, size, i), qelem(base, size, j), size);
	memcpy(qelem(base, size, j), tmp, size);
}


void
qreverse(void * base, size_t nmemb, size_t size)
{
	size_t i, j;

	if (nmemb < 2)
		return;

#if 0	// This goes slower?!?
	if (nmemb < BUFSIZ / size) {
		char tmp[nmemb * size];
		memcpy(tmp, base, nmemb * size);
		for (i = 0; i < nmemb; ++i)
			memcpy(base + i * size, tmp + (nmemb - 1 - i) * size, size);
		return;
	}
#endif

	for (i = 0, j = nmemb - 1; i < j; ++i, --j)
		qswap(base, size, i, j);
}


int
qpermute(void * base, size_t nmemb, size_t size,
		int(*compar)(const void *, const void *))
{
	size_t pivot, next = 0, i;

	if (nmemb < 2)
		// No other permutations!
		return 0;

	if (nmemb == 2) {
		if (compar(qelem(base, size, 0), qelem(base, size, 1)) >= 0)
			return 0;

		qswap(base, size, 0, 1);
		return 1;
	}

	for (pivot = nmemb - 1; pivot > 0; --pivot)
		if (compar(qelem(base, size, pivot - 1), qelem(base, size, pivot)) < 0)
			break;
	if (pivot == 0)
		// In final permutation (anti-sorted) order. None more remain.
		return 0;
	--pivot;

	// Find the next item to go in pivot's place, the right-most element
	// larger than pivot
	for (i = pivot + 1; i < nmemb; ++i)
		if (compar(qelem(base, size, pivot), qelem(base, size, i)) < 0)
			next = i;
	if (!next)
		return 0;

	// Swap the pivot with its replacement
	qswap(base, size, pivot, next);

	// Reverse all the elements after the new pivot
	qreverse(base + (pivot + 1) * size, nmemb - (pivot + 1), size);

	return 1;
}


int
ptrcmp(void const * a, void const * b)
{
	return *((void **) a) - *((void **) b);
}


struct valve {
	char name[3];
	int flow;
	int ndests;
	char * destnames;
	struct valve ** dests;
};


int
valve_ctor(struct valve * v, char const * name, int flow, char const * dests)
{
	int ndests;

	if (strlen(name) != 2
			|| (ndests = strlen(dests)) % 4 != 2)
		return -1;

	strcpy(v->name, name);
	v->flow = flow;
	v->ndests = (ndests + 2) / 4;
	v->destnames = strdup(dests);
	v->dests = calloc(ndests + 1, sizeof(struct valve *));

	return 0;
}


void
valve_dtor(struct valve * v)
{
	if (!v)
		return;
	free(v->destnames);
	free(v->dests);
}


void
valve_setdests(struct valve * v, struct valve * others, int nothers)
{
	int i, j;

	for (i = 0; i < v->ndests; ++i) {
		for (j = 0; j < nothers; ++j) {
			if (strncmp(v->destnames + i * 4, others[j].name, 2) == 0) {
				v->dests[i] = others + j;
				break;
			}
		}

		if (!v->dests[i])
			fprintf(stderr, "Cannot find %2.2s for %s\n", v->destnames + i * 4, v->name);
	}
}


void
valve_distances_dump(struct valve const * valves, int nvalves, int * dists)
{
	int i, j;

	fprintf(stderr, "   ");
	for (i = 0; i < nvalves; ++i) {
		fprintf(stderr, " %s", valves[i].name);
	}
	fprintf(stderr, "\n");

	for (i = 0; i < nvalves; ++i) {
		fprintf(stderr, " %s", valves[i].name);
		for (j = 0; j < nvalves; ++j) {
			fprintf(stderr, " %2d", dists[j + i * nvalves]);
		}
		fprintf(stderr, "\n");
	}
	fprintf(stderr, "\n");
}


// Precalculate distances between all pairs of valves
int *
valve_distances(struct valve * valves, int nvalves)
{
	int * dists;
	int i, j, n, more;
	struct valve ** v;

	if ((dists = malloc(nvalves * nvalves * sizeof(int))) == NULL)
		return NULL;

	// Initialise.
	for (i = 0; i < nvalves; ++i) {
		for (j = 0; j < nvalves; ++j) {
			dists[j + i * nvalves] = (i == j) ? 0 : -1;
		}
	}

	// Find successive distances
	for (n = 0, more = 1; more == 1; ++n) {
		more = 0;
		for (i = 0; i < nvalves; ++i) {
			// On each row.
			for (j = 0; j < nvalves; ++j) {
				// Look for valves that takes n steps to get to.
				if (dists[j + i * nvalves] != n)
					continue;

				for (v = valves[j].dests; v < valves[j].dests + valves[j].ndests; ++v) {
					if (dists[*v - valves + i * nvalves] != -1)
						continue;

					// Found a neighbor that we can't get to quicker
					// so we can get to it in n + 1 steps
					dists[(*v - valves) + i * nvalves] = n + 1;
					more = 1;
				}
			}
		}
	}

	return dists;
}


int
valve_distance(struct valve const * valves, int nvalves, int const * dists,
		struct valve const * a, struct valve const * b)
{
	return dists[(a - valves) + nvalves * (b - valves)];
}


// Calculate flow of a given path
int
path_flow(int * valves_reached,
		struct valve const * start, struct valve ** path, int pathlen,
		struct valve const * valves, int nvalves, int const * dists,
		int minutes)
{
	int flow = 0, i;

	for (i = 0; i < pathlen; ++i) {
		// Subtract travel time, and time to turn valve on
		minutes -= valve_distance(valves, nvalves, dists,
				i == 0 ? start : path[i - 1], path[i]);
		minutes -= 1;

		// Add the amount of flow for being turned on for remaining minutes
		if (minutes <= 0)
			break;

		flow += path[i]->flow * minutes;
	}

	*valves_reached += i;

	return flow;
}


// Calculate the best flow for a path, with a given set of plumbers
//
// Note that we only try cases where each subsequent plumber visits an equal or
// decreasing number of valves. This is because, given e.g.
//
//   a b c d e f g h
//
// If there are two plumbers with plenty of time to spare, the first might be
// able to get all the way up to 'f' by themselves. And we want to test all
// possible divisions of work to find the one with the most flow. But, we
// don't need to test the case where one plumber checks 'a' to 'c' and the
// other checks 'd' to 'h' here, because that will automatically get tested
// in the permutation.
//
//   d e f g h a b c
//
// So we only need to test equal-or-decreasing valves because increasing numbers
// of valves will be tested in a different iteration. Note that we still see
// some duplication of testing in the case of 'a' to 'd' and 'e' to 'h', but
// I can't see a way around that for now.
int
path_flow_plumbers(int * valves_reached,
		struct valve const * start, struct valve ** path, int pathlen,
		struct valve const * valves, int nvalves, int const * dists, int minutes,
		int plumbers, int max_valves)
{
	int flow_max = 0, flow_max_valves = 0;

	if (plumbers == 0)
		// No plumbers, no flow.
		return 0;

	// If there are 3 other plumbers, we don't want to try to reach the last
	// 3 valves, because we want to leave one each for them.
	// Or, if there are no other plumbers, we don't want to try more than
	// pathlen valves.
	if (max_valves > pathlen - (plumbers - 1))
		max_valves = pathlen - (plumbers - 1);

	if (plumbers == 1)
		// One plumber, simple case.
		return path_flow(valves_reached,
					start, path, max_valves,
					valves, nvalves, dists, minutes);

	// Try the maximum number of valves we can reach first, and the flow
	// that other plumbers can reach from there, and then loop backwards
	// through the number of valves we try to see which one gives the max
	// flow.
	while (max_valves > 0) {
		int flow, valves_self = 0, valves_rest = 0;

		flow = path_flow(&valves_self,
					start, path, max_valves,
					valves, nvalves, dists, minutes);
		flow += path_flow_plumbers(&valves_rest,
				start, path + valves_self, pathlen - valves_self,
				valves, nvalves, dists, minutes,
				plumbers - 1, valves_self);

		if (flow > flow_max) {
			// New max.
			flow_max = flow;
			flow_max_valves = valves_self + valves_rest;
		}

		// Can't reach all the valves from here. No improvements await
		if (valves_self + valves_rest < pathlen)
			break;

		max_valves = valves_self - 1;
	}

	*valves_reached += flow_max_valves;
	return flow_max;
}


int
main(int argc, char ** argv)
{
	int debug = 0, eruption = 30, plumbers = 1, i;
	regex_t valvein;
	regmatch_t rmvalve[4];
	char buf[BUFSIZ];
	struct valve * valves = NULL, * v, * start = NULL, ** path;
	int nvalves = 0, nflows = 0, max = 0;
	int * dists;

	if (regcomp(&valvein, "Valve ([[:upper:]]{2}) has flow rate=([[:digit:]]+);"
				" tunnels? leads? to valves? (([[:upper:]]{2}(, )?)+)", REG_EXTENDED) != 0)
	{
		fprintf(stderr, "Bad regex\n");
		return -1;
	}

	while ((i = getopt(argc, argv, "dp:b:t:")) != -1) {
		switch (i) {
		case 'd':
			debug = 1;
			break;

		case 'p':
			switch (atoi(optarg)) {
			case '1':
				eruption = 30;
				plumbers = 1;
				break;

			case 2:
				eruption = 26;
				plumbers = 2;
				break;

			default:
				fprintf(stderr, "Unexpected part %s\n", optarg);
				return -1;
			}
			break;

		case 'b':
			plumbers = atoi(optarg);
			break;

		case 't':
			// 'time' is taken by time(3), hence 'eruption'
			eruption = atoi(optarg);
			break;

		default:
			regfree(&valvein);
			return -1;
		}
	}

	// Read in valve data
	while (fgets(buf, sizeof(buf), stdin)
			&& regexec(&valvein, buf, arrlen(rmvalve), rmvalve, 0) == 0)
	{
		void * p;

		if ((p = realloc(valves, ++nvalves * sizeof(struct valve))) == NULL) {
			fprintf(stderr, "Bad realloc(%zd)\n", nvalves * sizeof(struct valve));
			free(valves);
			return -1;
		}
		valves = p;
		v = valves + nvalves - 1;

		buf[rmvalve[1].rm_eo] = '\0';
		buf[rmvalve[2].rm_eo] = '\0';
		buf[rmvalve[3].rm_eo] = '\0';
		valve_ctor(v, buf + rmvalve[1].rm_so,
				atoi(buf + rmvalve[2].rm_so),
				buf + rmvalve[3].rm_so);
	}

	// Set up valve dests
	for (v = valves; v < valves + nvalves; ++v)
		valve_setdests(v, valves, nvalves);

	// Calculate distances between all valves
	dists = valve_distances(valves, nvalves);
	if (debug)
		valve_distances_dump(valves, nvalves, dists);

	// Find start position
	for (v = valves; v < valves + nvalves && !start; ++v) {
		if (strcmp(v->name, "AA") == 0)
			start = v;
	}
	if (!start) {
		fprintf(stderr, "Unable to find start valve in %d valves\n", nvalves);
		for (v = valves; v < valves + nvalves; ++v)
			valve_dtor(v);
		free(valves);
		regfree(&valvein);
		return -1;
	}

	// Create a path containing all the valves that can flow, which are
	// reachable from the start position before the eruption
	path = malloc(nvalves * sizeof(struct valve *));
	for (v = valves; v < valves + nvalves; ++v) {
		if (v->flow > 0 && valve_distance(valves, nvalves, dists, start, v) < eruption)
			path[nflows++] = v;
	}

	// Go through all permutations of path to find max flow
	qsort(path, nflows, sizeof(struct valve *), ptrcmp);
	do {
		int flow, reached = 0;

		flow = path_flow_plumbers(&reached, start, path, nflows,
				valves, nvalves, dists, eruption,
				plumbers, nflows);

		// Is it a new high?
		if (flow > max) {
			max = flow;
			if (debug) {
				int j;
				fprintf(stderr, "Found new max flow: %d: %s", max, start->name);
				for (j = 0; j < nflows; ++j)
					fprintf(stderr, "->%s", path[j]->name);
				fprintf(stderr, "\n");
			}
		}

		if (reached < nflows) {
			// We didn't get to the end of the path. We can skip all
			// the permutations of all the remaining elements in the
			// path, because they don't matter.
			// Remaining elements will be in their first possible
			// permutation, i.e. in sorted order. If we reverse this
			// they'll be in their last permutation, and qpermute()
			// will then pick the next path that is substantively
			// different from this.
			qreverse(path + reached + 1, nflows - (reached + 1), sizeof(struct valve *));
		}
	} while (qpermute(path, nflows, sizeof(struct valve *), ptrcmp));

	printf("Max flow: %d\n", max);

	// Done.
	free(path);
	free(dists);
	for (v = valves; v < valves + nvalves; ++v)
		valve_dtor(v);
	free(valves);
	regfree(&valvein);

	return 0;
}