Aoc2021 day12

aoc2021-day12.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6bca0c88",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-12T11:50:22.121303Z",
     "start_time": "2021-12-12T11:50:21.581222Z"
    }
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from copy import copy\n",
    "\n",
    "from networkx.drawing.nx_pydot import graphviz_layout\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "from tqdm import tqdm\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cf8be1ac",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-12T11:50:22.129732Z",
     "start_time": "2021-12-12T11:50:22.124615Z"
    }
   },
   "outputs": [],
   "source": [
    "edges = \"\"\"fs-end\n",
    "he-DX\n",
    "fs-he\n",
    "start-DX\n",
    "pj-DX\n",
    "end-zg\n",
    "zg-sl\n",
    "zg-pj\n",
    "pj-he\n",
    "RW-he\n",
    "fs-DX\n",
    "pj-RW\n",
    "zg-RW\n",
    "start-pj\n",
    "he-WI\n",
    "zg-he\n",
    "pj-fs\n",
    "start-RW\"\"\"\n",
    "\n",
    "# edges = \"\"\"start-A\n",
    "# start-b\n",
    "# A-c\n",
    "# A-b\n",
    "# b-d\n",
    "# A-end\n",
    "# b-end\"\"\"\n",
    "\n",
    "# pos = {'start': (50, 100.0),\n",
    "#  'c': (0, 50.0),\n",
    "#  'A': (33.0, 50.0),\n",
    "#  'b': (66.0, 50.0),\n",
    "#  'd': (100.0, 50.0),\n",
    "#  'end': (50.0, 0.0)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c9e14eac",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-12T11:50:22.144394Z",
     "start_time": "2021-12-12T11:50:22.133352Z"
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "from collections import defaultdict, Counter\n",
    "\n",
    "G = defaultdict(lambda: list(), {})\n",
    "for line in edges.splitlines():\n",
    "    (l, r) = line.strip().split(\"-\")\n",
    "    if l != 'end' and r != 'start': G[l].append(r)\n",
    "    if r != 'end' and l != 'start': G[r].append(l)\n",
    "\n",
    "def pred_part1(n, path):\n",
    "    return n.isupper() or n not in path \n",
    "\n",
    "def pred_part2(n, path):\n",
    "    return pred_part1(n,path) or max(Counter(filter(str.islower, path)).values()) == 1\n",
    "\n",
    "def getpaths(current, path, pred):\n",
    "    if current == 'end': \n",
    "        return [path]\n",
    "\n",
    "    res = []\n",
    "    for x in G[current]:\n",
    "        if pred(x, path):\n",
    "            res += getpaths(x, path + [x], pred)\n",
    "\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "60322098",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-12T11:50:22.277990Z",
     "start_time": "2021-12-12T11:50:22.145910Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3509"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "paths2 = getpaths(\"start\", [\"start\"], pred_part2)\n",
    "len(paths2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f6d64899",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-12T11:50:22.284633Z",
     "start_time": "2021-12-12T11:50:22.279874Z"
    }
   },
   "outputs": [],
   "source": [
    "def graph_to_nx(g):\n",
    "    G = nx.DiGraph()\n",
    "    \n",
    "    for a in g:\n",
    "        G.add_node(a)\n",
    "        for b in g[a]:\n",
    "            G.add_edge(a, b, color='white', weight=10)\n",
    "    \n",
    "    return G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bc3553c",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2021-12-12T11:50:21.580Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 99%|█████████████████████████████▋| 99/100 [04:33<00:03,  3.27s/it]"
     ]
    }
   ],
   "source": [
    "G = graph_to_nx(G)\n",
    "pos = graphviz_layout(G, prog=\"circo\")\n",
    "\n",
    "def get_color(node, path, i):\n",
    "    if node == path[i]:\n",
    "        return 'green'\n",
    "    if node.isupper():\n",
    "        return 'blue'\n",
    "    if node.islower() and node in path[:i]:\n",
    "        return 'red'\n",
    "\n",
    "    return 'blue'\n",
    "\n",
    "def get_edge_color(path, i, u, v):\n",
    "    if (u, v) in list(zip(path[:i+1], path[1:i+1])) or (v, u) in list(zip(path[:i+1], path[1:i+1])):\n",
    "        return 'yellow'\n",
    "    return 'white'\n",
    "\n",
    "paths = paths2[:100]\n",
    "loop = tqdm(enumerate(paths), total=len(paths))\n",
    "cnt = 0\n",
    "for j, path in loop:\n",
    "    for i, p in enumerate(path):\n",
    "        fig, ax = plt.subplots(figsize=(10, 10))\n",
    "        ax.axis(\"off\")\n",
    "\n",
    "        edges = G.edges()\n",
    "        edge_color = [get_edge_color(path, i, u, v) for u,v in edges]\n",
    "        width = [G[u][v]['weight'] for u,v in edges]\n",
    "        node_color = [get_color(node, path, i) for node in G]\n",
    "\n",
    "        nx.draw_networkx(G, pos, ax=ax, with_labels=True, arrowsize=10, node_size=1000, node_color=node_color, edge_color=edge_color, width=width)\n",
    "\n",
    "        plt.savefig(f'output12/result_{cnt:05d}.svg')\n",
    "        plt.close()\n",
    "        cnt += 1\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e3c4bb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}