{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Binary Search Trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TreeNode:\n",
    "    def __init__(self, key, val, left=None, right=None, parent=None):\n",
    "        self.key = key\n",
    "        self.payload = val\n",
    "        self.left_child = left\n",
    "        self.right_child= right\n",
    "        self.parent = parent\n",
    "        \n",
    "    def __iter__(self):\n",
    "        if self:\n",
    "            if self.has_left_child():\n",
    "                for elem in self.left_child:\n",
    "                    yield elem\n",
    "                    \n",
    "            yield self.key\n",
    "            \n",
    "            if self.has_right_child():\n",
    "                for elem in self.right_child:\n",
    "                    yield elem\n",
    "                    \n",
    "    def has_left_child(self):\n",
    "        return self.left_child\n",
    "    \n",
    "    def has_right_child(self):\n",
    "        return self.right_child\n",
    "    \n",
    "    def is_left_child(self):\n",
    "        return self.parent and self.parent.left_child is self\n",
    "    \n",
    "    def is_right_child(self):\n",
    "        return self.parent and self.parent.right_child is self\n",
    "    \n",
    "    def is_root(self):\n",
    "        return not self.parent\n",
    "    \n",
    "    def is_leaf(self):\n",
    "        return not self.left_child and not self.right_child\n",
    "    \n",
    "    def has_any_children(self):\n",
    "        return self.right_child or self.left_child\n",
    "    \n",
    "    def has_both_children(self):\n",
    "        return self.right_child and self.left_child\n",
    "    \n",
    "    def replace_node_data(self, key, value, lc, rc):\n",
    "        self.key = key\n",
    "        self.payload = value\n",
    "        self.left_child = lc\n",
    "        self.right_child = rc\n",
    "        \n",
    "        if self.has_left_child():\n",
    "            self.left_child.parent = self\n",
    "            \n",
    "        if self.has_right_child():\n",
    "            self.right_child.parent = self\n",
    "            \n",
    "    def splice_out(self):\n",
    "        if self.is_leaf():\n",
    "            if self.is_left_child():\n",
    "                self.parent.left_child = None\n",
    "            else:\n",
    "                self.parent.right_child = None\n",
    "        elif self.has_any_children():\n",
    "            if self.has_left_child():\n",
    "                if self.is_left_child():\n",
    "                    self.parent.left_child = self.left_child\n",
    "                else:\n",
    "                    self.parent.right_child = self.left_child\n",
    "                    \n",
    "                self.left_child.parent = self.parent\n",
    "            else:  # has a right child\n",
    "                if self.is_left_child():\n",
    "                    self.parent.left_child = self.right_child\n",
    "                else:\n",
    "                    self.parent.right_child = self.right_child\n",
    "                    \n",
    "                self.right_child.parent = self.parent\n",
    "                \n",
    "    def find_successor(self):\n",
    "        succ = None\n",
    "        \n",
    "        if self.has_right_child():\n",
    "            succ = self.right_child.find_min()\n",
    "        else:  # no right child? Then we work on left child\n",
    "            if self.parent:  # it is an intermediate node\n",
    "                if self.is_left_child():\n",
    "                    succ = self.parent\n",
    "                else:  # is a right child\n",
    "                    self.parent.right_child = None\n",
    "                    succ = self.parent.find_successor()\n",
    "                    self.parent.right_child = self\n",
    "        return succ\n",
    "    \n",
    "    \n",
    "    def find_min(self):\n",
    "        current = self\n",
    "        \n",
    "        while current.has_left_child():\n",
    "            current = current.left_child\n",
    "\n",
    "        return current"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BinarySearchTree:\n",
    "    def __init__(self):\n",
    "        self.root = None\n",
    "        self.size = 0\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.length()\n",
    "    \n",
    "    def __iter__(self):\n",
    "        return self.root.__iter__()\n",
    "    \n",
    "    def __setitem__(self, k, v):\n",
    "        self.put(k, v)\n",
    "  \n",
    "    def __getitem__(self, k):\n",
    "        return self.get(k)\n",
    "    \n",
    "    def __contains__(self, k):\n",
    "        #if self._get(key, self.root):\n",
    "        #    return True\n",
    "        #else:\n",
    "        #    return False\n",
    "        \n",
    "        return self._get(k, self.root) is not None\n",
    "\n",
    "    def __delitem__(self, k):\n",
    "        self.delete(k)\n",
    "    \n",
    "    def length(self):\n",
    "        return self.size\n",
    "    \n",
    "    def put(self, key, val):\n",
    "        if self.root:\n",
    "            self._put(key, val, self.root)\n",
    "        else:  # the tree is non-existient\n",
    "            self.root = TreeNode(key, val)\n",
    "            \n",
    "        self.size += 1\n",
    "        \n",
    "    def _put(self, key, val, current_node):\n",
    "        if key < current_node.key:\n",
    "            if current_node.has_left_child():\n",
    "                self._put(key, val, current_node.left_child)\n",
    "            else:\n",
    "                current_node.left_child = TreeNode(key, val, parent=current_node)\n",
    "        else:\n",
    "            if current_node.has_right_child():\n",
    "                self._put(key, val, current_node.right_child)\n",
    "            else:\n",
    "                current_node.right_child = TreeNode(key, val, parent=current_node)\n",
    "                \n",
    "    def get(self, key):\n",
    "        if self.root:\n",
    "            res = self._get(key, self.root)\n",
    "            if res:\n",
    "                # return the value of our TreeNode with matching key\n",
    "                return res.payload\n",
    "            else:\n",
    "                return None\n",
    "        else:\n",
    "            return None\n",
    "        \n",
    "    def _get(self, key, current_node):\n",
    "        if not current_node:\n",
    "            return None\n",
    "        elif current_node.key == key:\n",
    "            return current_node\n",
    "        elif key < current_node.key:\n",
    "            return self._get(key, current_node.left_child)\n",
    "        else:\n",
    "            return self._get(key, current_node.right_child)\n",
    "        \n",
    "    def delete(self, key):\n",
    "        if self.size > 1:\n",
    "            node_to_remove = self._get(key, self.root)\n",
    "            if node_to_remove:\n",
    "                self.remove(node_to_remove)\n",
    "                self.size -= 1\n",
    "            else:\n",
    "                raise KeyError(\"Key not found in tree!\")\n",
    "        elif self.size == 1 and self.root.key == key:\n",
    "            # tree is completely removed, since it was only one node\n",
    "            self.root = None\n",
    "            self.size -= 1\n",
    "        else:\n",
    "            raise KeyError(\"Key not found in tree!\")\n",
    "            \n",
    "    def remove(self, current_node):\n",
    "        if current_node.is_leaf():  # leaf\n",
    "            #if current_node == current_node.parent.left_child:\n",
    "            if current_node.is_left_child():   \n",
    "                current_node.parent.left_child = None\n",
    "            else:\n",
    "                current_node.parent.right_child = None\n",
    "        elif current_node.has_both_children():  # interior or root\n",
    "            succ = current_node.find_successor()\n",
    "            succ.splice_out()\n",
    "            current_node.key = succ.key\n",
    "            current_node.payload = succ.payload\n",
    "        else:  # node has one child\n",
    "            if current_node.has_left_child():\n",
    "                if current_node.is_left_child():\n",
    "                    current_node.left_child.parent = current_node.parent\n",
    "                    current_node.parent.left_child = current_node.left_child\n",
    "                elif current_node.is_right_child():\n",
    "                    current_node.left_child.parent = current_node.parent\n",
    "                    current_node.parent.right_child = current_node.left_child\n",
    "                else: #root\n",
    "                    current_node.replace_node_data(current_node.left_child.key,\n",
    "                                                   current_node.left_child.payload,\n",
    "                                                   current_node.left_child.left_child,\n",
    "                                                   current_node.left_child.right_child)\n",
    "            else: # have right child\n",
    "                if current_node.is_left_child():\n",
    "                    current_node.right_child.parent = current_node.parent\n",
    "                    current_node.parent.left_child = current_node.right_child\n",
    "                elif current_node.is_right_child():\n",
    "                    current_node.right_child.parent = current_node.parent\n",
    "                    current_node.parent.right_child = current_node.right_child\n",
    "                else:\n",
    "                    current_node.replace_node_data(current_node.right_child.key,\n",
    "                                                   current_node.right_child.payload,\n",
    "                                                   current_node.right_child.left_child,\n",
    "                                                   current_node.right_child.right_child)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing Our Binary Search Tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "bt = BinarySearchTree()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.BinarySearchTree at 0x7f4328211b38>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "bt[24] = \"apple\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'apple'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt[24]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'apple'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt.get(24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "24 in bt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "42 in bt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(bt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.TreeNode at 0x7f432820d550>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt.root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "del bt[24]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(bt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Don't Panic!\""
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt[42]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "bt[42] = \"Don't Panic!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "'Key not found in tree!'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-24-72e1e3b14341>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mbt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdelete\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m24\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-4-58f4d14e8ee9>\u001b[0m in \u001b[0;36mdelete\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m     78\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m \u001b[0;34m-=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     79\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 80\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mKeyError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Key not found in tree!\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     81\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mroot\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkey\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     82\u001b[0m             \u001b[0;31m# tree is completely removed, since it was only one node\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'Key not found in tree!'"
     ]
    }
   ],
   "source": [
    "bt.delete(24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "bt[24] = \"Panic Don't!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[24, 42, 42, 42, 42, 42, 42]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(bt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = {42: \"Don't Panic!\", 24: \"Panic Don't!\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Panic Don't!\""
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d[24]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Panic Don't!\""
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt[24]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{42: \"Don't Panic!\", 24: \"Panic Don't!\"}"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.BinarySearchTree at 0x7f4328211b38>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "24: \"Panic Don't!\"\n",
      "42: \"Don't Panic!\"\n",
      "42: \"Don't Panic!\"\n",
      "42: \"Don't Panic!\"\n",
      "42: \"Don't Panic!\"\n",
      "42: \"Don't Panic!\"\n",
      "42: \"Don't Panic!\"\n",
      "42: \"Don't Panic!\"\n"
     ]
    }
   ],
   "source": [
    "for k in bt:\n",
    "    print(f\"{k}: {bt[k]!r}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing with names and ranks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from random import shuffle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_names(url):\n",
    "    req = requests.get(url)\n",
    "    lines = req.text.split(\"\\n\")\n",
    "    names = []\n",
    "    \n",
    "    for line in lines:\n",
    "        data = line.split()\n",
    "        if len(data) == 4:\n",
    "            names.append( (data[0], data[3]) )\n",
    "            \n",
    "    return names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    names = get_names(\"https://www2.census.gov/topics/genealogy/1990surnames/dist.female.first\")\n",
    "    print(f\"Extracted {len(names)} names\")\n",
    "    shuffle(names)\n",
    "    \n",
    "    name_tree = BinarySearchTree()\n",
    "    \n",
    "    for name, rank in names:\n",
    "        name_tree[int(rank)] = name\n",
    "        \n",
    "    return name_tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracted 4275 names\n"
     ]
    }
   ],
   "source": [
    "nt = main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4275"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(nt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ANNA'"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nt[33]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'MARY'"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nt[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'YURIKO'"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nt[4000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'CELINA'"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nt[1000]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing our BST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphviz import Digraph\n",
    "\n",
    "class Stack:\n",
    "    def __init__(self):\n",
    "         self.items = []\n",
    "\n",
    "    def is_empty(self):\n",
    "         return self.items == []\n",
    "\n",
    "    def push(self, item):\n",
    "        self.items.append(item)\n",
    "\n",
    "    def pop(self):\n",
    "         return self.items.pop()\n",
    "\n",
    "    def peek(self):\n",
    "         return self.items[-1]\n",
    "\n",
    "    def size(self):\n",
    "         return len(self.items)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viz_tree(r):\n",
    "    stack = Stack()\n",
    "    g = Digraph(node_attr={'shape': 'record', 'height': '.1'})\n",
    "    _id = 0\n",
    "    \n",
    "#    if isinstance(r, BinaryTree):\n",
    "#        current_node = r\n",
    "#    elif isinstance(r, BinarySearchTree):\n",
    "#        current_node = r.root  # root is a TreeNode object!    \n",
    "    current_node = r.root\n",
    "    \n",
    "    leftward = True\n",
    "    current_root_num = 0\n",
    "\n",
    "    while True:\n",
    "        if current_node:\n",
    "            stack.push((_id, current_node))\n",
    "            \n",
    "            if isinstance(current_node, TreeNode):\n",
    "                g.node(f'node{_id}',\n",
    "                   f'<f0>|<f1> {current_node.key}:{current_node.payload}|<f2> ')\n",
    "            elif isinstance(current_node, AVLTreeNode):\n",
    "                g.node(f'node{_id}',\n",
    "                       f'<f0>|<f1> {current_node.key}:{current_node.payload} ({current_node.balance_factor})|<f2> ')\n",
    "\n",
    "            if _id >= 1:\n",
    "                g.edge(f'node{current_root_num}:f{0 if leftward else 2}',\n",
    "                       f'node{_id}:f1')\n",
    "\n",
    "            leftward = True\n",
    "            current_node = current_node.left_child  # left\n",
    "            current_root_num = _id\n",
    "            _id += 1\n",
    "\n",
    "        if current_node is None and not stack.is_empty():\n",
    "            count, popped_node = stack.pop()\n",
    "            if popped_node.right_child:\n",
    "                current_root_num = count\n",
    "                current_node = popped_node.right_child  # right\n",
    "                leftward = False\n",
    "\n",
    "        if current_node is None and stack.is_empty():\n",
    "            break\n",
    "\n",
    "    return g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "g = viz_tree(nt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'bst.pdf'"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g.render(\"bst\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing with 'names' as keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    names = get_names(\"https://www2.census.gov/topics/genealogy/1990surnames/dist.female.first\")\n",
    "    print(f\"Extracted {len(names)} names\")\n",
    "    shuffle(names)\n",
    "    \n",
    "    name_tree = BinarySearchTree()\n",
    "    \n",
    "    for name, rank in names:\n",
    "        name_tree[name] = int(rank)\n",
    "        \n",
    "    return name_tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = viz_tree(nt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.render(\"bst_names_as_key\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing with sorted 'names' as keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    names = get_names(\"https://www2.census.gov/topics/genealogy/1990surnames/dist.female.first\")\n",
    "    print(f\"Extracted {len(names)} names\")\n",
    "    \n",
    "    # No shuffling!\n",
    "    #shuffle(names)\n",
    "    \n",
    "    # names[(\"MARY\", \"1\"), ...)\n",
    "    names.sort(key=lambda t: t[0])\n",
    "    \n",
    "    name_tree = BinarySearchTree()\n",
    "    \n",
    "    for name, rank in names:\n",
    "        name_tree[name] = int(rank)\n",
    "        \n",
    "    return name_tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = viz_tree(nt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.render(\"bst_names_as_key_sorted\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing with sorted 'names' as keys, smaller subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    names = get_names(\"https://www2.census.gov/topics/genealogy/1990surnames/dist.female.first\")\n",
    "    print(f\"Extracted {len(names)} names\")\n",
    "    \n",
    "    # No shuffling!\n",
    "    #shuffle(names)\n",
    "    \n",
    "    # names[(\"MARY\", \"1\"), ...)\n",
    "    names.sort(key=lambda t: t[0])\n",
    "    \n",
    "    name_tree = BinarySearchTree()\n",
    "    \n",
    "    for name, rank in names[:100]:  # limit our tree to 100 names!\n",
    "        name_tree[name] = int(rank)\n",
    "        \n",
    "    return name_tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = viz_tree(nt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.render(\"bst_names_as_key_sorted\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
