recursive tree

Author: zedrdave

2018/Playground.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": 117,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -15,7 +15,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 52,
+   "execution_count": 101,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -24,7 +24,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 78,
+   "execution_count": 125,
    "metadata": {
     "scrolled": false
    },
@@ -58,6 +58,60 @@
     "    print(f\"Part {p}:\", recur(map(int, data.split()), part=p)[1])"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 126,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Part 1: 36566\n",
+      "Part 2: 30548\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Approach 2\n",
+    "# \n",
+    "# Recursively build tree first\n",
+    "\n",
+    "from dataclasses import dataclass\n",
+    "from typing import List, Union\n",
+    "\n",
+    "@dataclass\n",
+    "class node:\n",
+    "    meta: int = 0\n",
+    "    children: List['node'] = ()\n",
+    "        \n",
+    "    def part_1(self) -> int:\n",
+    "        return sum(self.meta) + sum(n.part_1() for n in self.children)\n",
+    "\n",
+    "    def part_2(self) -> int:\n",
+    "        if self.children:\n",
+    "            return sum(self.children[m-1].part_2() \n",
+    "                       for m in self.meta if 0 < m <= len(self.children))\n",
+    "        else:\n",
+    "            return sum(self.meta)\n",
+    "\n",
+    "        \n",
+    "def recur_tree(N): \n",
+    "    n_children, n_meta, *N = N\n",
+    "    \n",
+    "    children = []\n",
+    "    for _ in range(n_children):\n",
+    "        N, child = recur_tree(N)\n",
+    "        children += [child]\n",
+    "\n",
+    "    return N[n_meta:], node(meta=N[:n_meta], children=children)\n",
+    "    \n",
+    "tree = recur_tree(map(int, data.split()))[1]\n",
+    "\n",
+    "print(f\"Part 1:\", tree.part_1())\n",
+    "print(f\"Part 2:\", tree.part_2())"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,