2022/Day 19.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 19, robotic pathfinding\n",
    "\n",
    "This year's prevailing puzzle theme appears to be pathfinding! This is very much like [this year's day 16 puzzle](./Day%2016.ipynb), in that there is a limited time to 'open' resource-producing robots, and crack open the most geodes.\n",
    "\n",
    "Instead of you moving around a set of tunnels and opening valves one by one, we have robots bring resources to a single robot factory, so that's going to be the focus of a BFS traversal. The factory needs resources produced by robots, but it can only use a certain amount of those resources when it works, putting an upper limit on the number of robots you'd want to output (e.g. in the first blueprint in the example, the factory will need, at most, 4 ore for a robot type, so we can avoid wasting time building more than 3 additional ore-producing robots to add to your starter robot).\n",
    "\n",
    "One thing to realise is that there is a theoretical upper limit on the number of geodes that can be cracked in the time given: if you had the resources to produce a new geode-cracking robot every minute, then the maximum number of geodes produced by the robots that'll be built in time $T$ is the [triangle number of $T - 1$](https://en.wikipedia.org/wiki/Triangular_number), so $\\frac {T(T - 1)} {2}$ (the robot you produce _this_ minute starts outputing geodes the next minute, hence $T - 1$). That property can be used to prune the search tree further; if a given state will only produce at most _max geodes cracked with existing robots + geode-cracking potential_, and we already have found another state that cracks more geodes, we can discard this specific path.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pyright: strict\n",
    "from collections import deque\n",
    "from dataclasses import dataclass\n",
    "from enum import IntEnum\n",
    "from typing import Iterable, Iterator, Literal, NamedTuple, Self, TypeAlias\n",
    "\n",
    "Amount: TypeAlias = int\n",
    "# The number of robots we have\n",
    "Robots: TypeAlias = tuple[Amount, Amount, Amount, Amount]\n",
    "# The number of resources we can build in the time, with the robots we have\n",
    "Resources: TypeAlias = tuple[Amount, Amount, Amount, Amount]\n",
    "# the amount of resources the factory requires to build one robot (the factory\n",
    "# only needs ore, clay and obsidian).\n",
    "Recipe: TypeAlias = tuple[Amount, Amount, Amount]\n",
    "\n",
    "\n",
    "class Resource(IntEnum):\n",
    "    ore = 0\n",
    "    clay = 1\n",
    "    obsidian = 2\n",
    "    geode = 3\n",
    "\n",
    "\n",
    "class RobotFactoryState(NamedTuple):\n",
    "    remaining: int\n",
    "    robots: Robots = (1, 0, 0, 0)\n",
    "    resources: Resources = (0, 0, 0, 0)\n",
    "\n",
    "    def traverse(self, bp: \"Blueprint\") -> Iterator[\"RobotFactoryState\"]:\n",
    "        rem, robots, resources = self.remaining, self.robots, self.resources\n",
    "        per_robot = zip(Resource, bp.recipes, robots, bp.max_robots, resources)\n",
    "        # what robot to build next?\n",
    "        for rtype, recipe, have, rmax, res in per_robot:\n",
    "            if rmax and have * rem + res >= rmax * rem:\n",
    "                # factory can't consume more of this resource, no point in\n",
    "                # producing this type.\n",
    "                continue\n",
    "            if not all(bool(prod) for prod, req in zip(robots, recipe) if req):\n",
    "                # Not all resources can be produced yet\n",
    "                continue\n",
    "            # how much time do we need to produce enough of each required resource?\n",
    "            needed = 1 + (\n",
    "                max(\n",
    "                    0 if avail >= req else (req - avail + prod - 1) // prod\n",
    "                    for req, avail, prod in zip(recipe, resources, robots)\n",
    "                    if req\n",
    "                )\n",
    "            )\n",
    "            if needed >= rem:  # no time left to build this robot\n",
    "                continue\n",
    "            # produce a new state, with the resources that'll be made available\n",
    "            # by the time the new robot is done minus the resources consumed by\n",
    "            # the factory to build the new robot, and the new number of robots\n",
    "            # with this type incremented.\n",
    "            new_res = [\n",
    "                avail - req + prod * needed\n",
    "                for avail, req, prod in zip(resources, (*recipe, 0), robots)\n",
    "            ]\n",
    "            new_robots = list(robots)\n",
    "            new_robots[rtype] += 1\n",
    "            yield RobotFactoryState(\n",
    "                rem - needed,\n",
    "                (new_robots[0], new_robots[1], new_robots[2], new_robots[3]),\n",
    "                (new_res[0], new_res[1], new_res[2], new_res[3]),\n",
    "            )\n",
    "\n",
    "    @property\n",
    "    def max_geodes(self) -> Amount:\n",
    "        \"\"\"Max geodes this state can produce given the remaining time and built robots\"\"\"\n",
    "        return (\n",
    "            self.resources[Resource.geode]\n",
    "            + self.remaining * self.robots[Resource.geode]\n",
    "        )\n",
    "\n",
    "    @property\n",
    "    def max_geode_potential(self) -> Amount:\n",
    "        \"\"\"Max geodes this state can produce given the remaining time, built robots, and potential robots\"\"\"\n",
    "        return self.max_geodes + self.remaining * (self.remaining - 1) // 2\n",
    "\n",
    "    @property\n",
    "    def priority(self) -> tuple[Amount, Amount, Amount, int]:\n",
    "        \"\"\"Priority queue key\n",
    "\n",
    "        Prioritises states by geodes produced, obsidian produced, clay produced and time remaining.\n",
    "\n",
    "        \"\"\"\n",
    "        _, clay, obsidian, geodes = self.resources\n",
    "        return (-geodes, -obsidian, -clay, self.remaining)\n",
    "\n",
    "\n",
    "@dataclass(frozen=True)\n",
    "class Blueprint:\n",
    "    # the recipe for each robot type\n",
    "    recipes: tuple[Recipe, Recipe, Recipe, Recipe]\n",
    "    # the upper limit for each robot, set by the maximum amount of each resource\n",
    "    # the factory can utilise. (The factory never needs Geode-cracking robots)\n",
    "    max_robots: tuple[Amount, Amount, Amount, Literal[0]]\n",
    "\n",
    "    @classmethod\n",
    "    def from_line(cls, line: str) -> Self:\n",
    "        # we only need the numbers, that happen to appear on indices that\n",
    "        # are multiples of 3.\n",
    "        words = line.split()[6::3]\n",
    "        recipes = (\n",
    "            # ore robot, amount of ore\n",
    "            (int(words[0]), 0, 0),\n",
    "            # clay robot, amount of ore\n",
    "            (int(words[2]), 0, 0),\n",
    "            # obsidian robot, amount of ore and clay\n",
    "            (int(words[4]), int(words[5]), 0),\n",
    "            # geode-cracking robot, amount of ore and obsidian\n",
    "            (int(words[7]), 0, int(words[8])),\n",
    "        )\n",
    "        max_bots = [max(resource) for resource in zip(*recipes)]\n",
    "        # no limit on the number of geode bots\n",
    "        return cls(recipes, (max_bots[0], max_bots[1], max_bots[2], 0))\n",
    "\n",
    "    def maximum_opened_geodes(self, time: int) -> int:\n",
    "        start = RobotFactoryState(time)\n",
    "        queue = deque([start])\n",
    "        seen = {start}\n",
    "        max_geodes = 0\n",
    "        while queue:\n",
    "            for state in queue.popleft().traverse(self):\n",
    "                if state in seen or state.max_geode_potential < max_geodes:\n",
    "                    continue\n",
    "                max_geodes = max(max_geodes, state.max_geodes)\n",
    "                seen.add(state)\n",
    "                queue.append(state)\n",
    "        return max_geodes\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Factory:\n",
    "    blueprints: list[Blueprint]\n",
    "\n",
    "    @classmethod\n",
    "    def from_lines(cls, lines: Iterable[str]) -> Self:\n",
    "        return cls([Blueprint.from_line(bp) for bp in lines])\n",
    "\n",
    "    def quality_levels(self, time: int = 24) -> Iterator[int]:\n",
    "        for i, bp in enumerate(self.blueprints, 1):\n",
    "            yield bp.maximum_opened_geodes(time) * i\n",
    "\n",
    "\n",
    "example = Factory.from_lines(\n",
    "    # multi-line blueprints still work as we split by variable-length whitespace.\n",
    "    \"\"\"\\\n",
    "Blueprint 1:\n",
    "  Each ore robot costs 4 ore.\n",
    "  Each clay robot costs 2 ore.\n",
    "  Each obsidian robot costs 3 ore and 14 clay.\n",
    "  Each geode robot costs 2 ore and 7 obsidian.\n",
    "\n",
    "Blueprint 2:\n",
    "  Each ore robot costs 2 ore.\n",
    "  Each clay robot costs 3 ore.\n",
    "  Each obsidian robot costs 3 ore and 8 clay.\n",
    "  Each geode robot costs 3 ore and 12 obsidian.\n",
    "\"\"\".split(\"\\n\\n\")\n",
    ")\n",
    "\n",
    "assert sum(example.quality_levels()) == 33"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 1: 1480\n"
     ]
    }
   ],
   "source": [
    "import aocd\n",
    "\n",
    "factory = Factory.from_lines(aocd.get_data(day=19, year=2022).splitlines())\n",
    "print(\"Part 1:\", sum(factory.quality_levels()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2, run for longer\n",
    "\n",
    "For part two, we need to extend the runtime of our simulation. Those 8 extra minutes produce a lot more states!\n",
    "\n",
    "I switched the BFS over to a priority queue, focusing on states with geodes before handling other states, as that helps produce a maximum geode count sooner and so helps prune unproductive branches faster. This reduces the execution time for part 2 from ~3 seconds down to 1.6 seconds. It has an even greater effect on the example blueprints; there the time goes from well over 6 seconds down to about 2 seconds.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import reduce\n",
    "from heapq import heappop, heappush\n",
    "from itertools import count\n",
    "from operator import mul\n",
    "\n",
    "\n",
    "def prioritised_maximum_opened_geodes(bp: Blueprint) -> int:\n",
    "    tiebreaker = count()\n",
    "    queue: list[tuple[Amount, Amount, Amount, int, int, RobotFactoryState]] = []\n",
    "\n",
    "    def add(state: RobotFactoryState):\n",
    "        heappush(queue, (*state.priority, next(tiebreaker), state))\n",
    "\n",
    "    start = RobotFactoryState(32)\n",
    "    add(start)\n",
    "    seen = {start}\n",
    "    max_geodes = 0\n",
    "    while queue:\n",
    "        *_, state = heappop(queue)\n",
    "        for nstate in state.traverse(bp):\n",
    "            if nstate in seen or nstate.max_geode_potential < max_geodes:\n",
    "                continue\n",
    "            max_geodes = max(max_geodes, nstate.max_geodes)\n",
    "            seen.add(nstate)\n",
    "            add(nstate)\n",
    "    return max_geodes\n",
    "\n",
    "\n",
    "def largest_number_geodes(factory: Factory) -> Iterator[int]:\n",
    "    for bp in factory.blueprints[:3]:\n",
    "        yield prioritised_maximum_opened_geodes(bp)\n",
    "\n",
    "\n",
    "assert reduce(mul, largest_number_geodes(example)) == 3472"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 2: 3168\n"
     ]
    }
   ],
   "source": [
    "print(\"Part 2:\", reduce(mul, largest_number_geodes(factory)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "adventofcode-mOkh6lsX",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b1b6870d1e0a983b1943c858d70ac8a7c80477f9f3ca364eb8daa198319a8a87"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}