Aoc2021 day06 lanternfish got out of hand

aoc2021-day06.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4d64fdb8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-06T21:05:46.912768Z",
     "start_time": "2021-12-06T21:05:46.267963Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import PIL\n",
    "import pydot\n",
    "\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm\n",
    "from networkx.drawing.nx_pydot import graphviz_layout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c17761c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-06T21:05:46.922532Z",
     "start_time": "2021-12-06T21:05:46.913910Z"
    }
   },
   "outputs": [],
   "source": [
    "fish = [\n",
    "    1, 1, 1, 1, 1, 1, 1, 4, 1, 2, 1, 1, 4, 1, 1, 1, 5, 1, 1, 1, 1, 1, 1, 1, 1,\n",
    "    1, 1, 1, 1, 5, 1, 1, 1, 1, 3, 1, 1, 2, 1, 2, 1, 3, 3, 4, 1, 4, 1, 1, 3, 1,\n",
    "    1, 5, 1, 1, 1, 1, 4, 1, 1, 5, 1, 1, 1, 4, 1, 5, 1, 1, 1, 3, 1, 1, 5, 3, 1,\n",
    "    1, 1, 1, 1, 4, 1, 1, 1, 1, 1, 2, 4, 1, 1, 1, 1, 4, 1, 2, 2, 1, 1, 1, 3, 1,\n",
    "    2, 5, 1, 4, 1, 1, 1, 3, 1, 1, 4, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 4, 1, 1, 1,\n",
    "    1, 1, 1, 1, 2, 1, 1, 5, 1, 1, 1, 4, 1, 1, 5, 1, 1, 5, 3, 3, 5, 3, 1, 1, 1,\n",
    "    4, 1, 1, 1, 1, 1, 1, 5, 3, 1, 2, 1, 1, 1, 4, 1, 3, 1, 5, 1, 1, 2, 1, 1, 1,\n",
    "    1, 1, 5, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 4, 3, 2, 1, 2, 4, 1, 3, 1, 5, 1, 2,\n",
    "    1, 4, 1, 1, 1, 1, 1, 3, 1, 4, 1, 1, 1, 1, 3, 1, 3, 3, 1, 4, 3, 4, 1, 1, 1,\n",
    "    1, 5, 1, 3, 3, 2, 5, 3, 1, 1, 3, 1, 3, 1, 1, 1, 1, 4, 1, 1, 1, 1, 3, 1, 5,\n",
    "    1, 1, 1, 4, 4, 1, 1, 5, 5, 2, 4, 5, 1, 1, 1, 1, 5, 1, 1, 2, 1, 1, 1, 1, 1,\n",
    "    1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 5, 1, 1, 1, 1, 1, 1, 3, 1, 1, 2, 1, 1\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6883b1d5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-06T21:05:46.933583Z",
     "start_time": "2021-12-06T21:05:46.925671Z"
    }
   },
   "outputs": [],
   "source": [
    "class Fish:\n",
    "    def __init__(self, birth_turn, left=None, right=None):\n",
    "        self.left = left\n",
    "        self.right = right\n",
    "        self.birth_turn = birth_turn\n",
    "\n",
    "class FunSolution:\n",
    "    def solve(self, fish, n):\n",
    "        def handleFish(x, n):\n",
    "            if x < n:\n",
    "                fish_a = handleFish(x+7, n)\n",
    "                fish_b = handleFish(x+9, n)\n",
    "                return Fish(x, fish_a, fish_b)\n",
    "            else:\n",
    "                return Fish(x)\n",
    "    \n",
    "        return handleFish(fish, n)\n",
    "    \n",
    "def bfs_fish_to_nx(parent):\n",
    "    G = nx.DiGraph()\n",
    "    q = [parent]\n",
    "    \n",
    "    while q:\n",
    "        parent = q.pop()\n",
    "        \n",
    "        G.add_node(parent)\n",
    "        if parent.left == None or parent.right == None:\n",
    "            continue\n",
    "        \n",
    "        G.add_edge(parent, parent.left, color='black', weight=1)\n",
    "        G.add_edge(parent, parent.right, color='green', weight=1)\n",
    "        \n",
    "        q.extend([parent.left, parent.right])\n",
    "        \n",
    "    return G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c348c7b8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-06T21:21:11.823959Z",
     "start_time": "2021-12-06T21:14:24.313808Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 5/5 [06:47<00:00, 81.50s/it]\n"
     ]
    }
   ],
   "source": [
    "fish_img = PIL.Image.open('icons/fish.png')\n",
    "\n",
    "for n_days in tqdm(range(0, 80)):\n",
    "    Fish.fish_count = 0\n",
    "    parent = FunSolution().solve(fish[0], n_days)\n",
    "    G = bfs_fish_to_nx(parent)\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(20, 20))\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "    edges = G.edges()\n",
    "    colors = [G[u][v]['color'] for u,v in edges]\n",
    "    weights = [G[u][v]['weight']/2 for u,v in edges]\n",
    "    pos = graphviz_layout(G, prog=\"dot\")    \n",
    "    pos = {k: (v[0]*2, -1 * min(k.birth_turn, n_days+1)) for k, v in pos.items()}\n",
    "    nx.draw_networkx(G, pos, ax=ax, with_labels=False, node_color='w', arrowsize=10, edge_color=colors, width=weights)\n",
    "\n",
    "    tr_figure = ax.transData.transform\n",
    "    tr_axes = fig.transFigure.inverted().transform\n",
    "\n",
    "    icon_size = np.log(plt.xlim()[1]) / n_days / 10\n",
    "    icon_center = icon_size / 2.0\n",
    "    for n in G.nodes:\n",
    "        if n.birth_turn > n_days:\n",
    "            continue\n",
    "        xf, yf = tr_figure(pos[n])\n",
    "        xa, ya = tr_axes((xf, yf))\n",
    "        a = plt.axes([xa - icon_center, ya - icon_center, icon_size, icon_size])\n",
    "        a.imshow(fish_img)\n",
    "        a.axis(\"off\")\n",
    "\n",
    "    plt.savefig(f'output/result{n_days:05d}.png')\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe8f7bc3",
   "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
}