Chapter 3.5 --- Red-black trees.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Red-black trees\n",
    "\n",
    "* Red-black trees are a type of self-balancing binary search tree. They have $O(\\log{n})$ read times, making them highly efficient.\n",
    "* In addition to the properties of a binary search tree red-black trees have the following invariants:\n",
    "  * Each node is either red or black.\n",
    "  * All leaves are `null` and black.\n",
    "  * If a node is red, then both its children are black.\n",
    "  * Every path from a given node to any of its descendant `null` nodes contains the same number of black nodes.\n",
    "* Red–black trees offer worst-case guarantees for insertion time, deletion time, and search time.\n",
    "* Read operations are the same as in binary search trees.\n",
    "* Insert and delete are much more compliated, but still $O(\\log{n})$.\n",
    "* Here's the tutorial I'm following: http://staff.ustc.edu.cn/~csli/graduate/algorithms/book6/chap14.htm.\n",
    "\n",
    "\n",
    "* A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node:\n",
    "    def __init__(self, color, left, right, parent, value=None):\n",
    "        self.color = color\n",
    "        self.left = left\n",
    "        self.right = right\n",
    "        self.parent = parent\n",
    "        self.value = value\n",
    "        \n",
    "\n",
    "class Tree:\n",
    "    def __init__(self, root):\n",
    "        self.root = root\n",
    "        \n",
    "        \n",
    "def set_left(node, parent):\n",
    "    # If the node is the right child of the parent, remove it from there.\n",
    "    if parent.right == node:\n",
    "        parent.right = None\n",
    "    \n",
    "    # If the parent has an existing left child set that child's parent to None (overwrite).\n",
    "    if parent.left is not None:\n",
    "        parent.left.parent = None\n",
    "    \n",
    "    # Attach the node to the parent.\n",
    "    parent.left = node\n",
    "    \n",
    "    # If the node is non-null assign it the parent pointer.\n",
    "    if node is not None:\n",
    "        node.parent = parent\n",
    "    \n",
    "        \n",
    "def set_right(node, parent):\n",
    "    if parent.left == node:\n",
    "        parent.left = None\n",
    "    \n",
    "    if parent.right is not None:\n",
    "        parent.right.parent = None    \n",
    "    \n",
    "    parent.right = node\n",
    "    \n",
    "    if node is not None:\n",
    "        node.parent = parent\n",
    "\n",
    "    \n",
    "def left_rotate(T, x):\n",
    "    y = x.right\n",
    "    \n",
    "    set_right(node=y.left, parent=x)\n",
    "    \n",
    "    if x.parent is None:\n",
    "        T.root = y  # TODO: ?\n",
    "    else:\n",
    "        if x == x.parent.left:\n",
    "            set_left(node=y, parent=x.parent)\n",
    "        else:\n",
    "            set_right(node=y, parent=x.parent)\n",
    "            \n",
    "    set_left(node=x, parent=y)\n",
    "    \n",
    "    \n",
    "def right_rotate(T, x):\n",
    "    y = x.left\n",
    "    \n",
    "    set_left(node=y.right, parent=x)\n",
    "    \n",
    "    if x.parent is None:\n",
    "        T.root = y\n",
    "    else:\n",
    "        if x == x.parent.right:\n",
    "            set_right(node=y, parent=x.parent)\n",
    "        else:\n",
    "            set_left(node=y, parent=x.parent)\n",
    "            \n",
    "    set_right(node=x, parent=y)\n",
    "    \n",
    "    \n",
    "def binary_insert(T, val):\n",
    "    \"\"\"Insert as though inserting into a binary tree, with the color red.\"\"\"\n",
    "    n = T.root\n",
    "    while True:\n",
    "        if val < n.value:\n",
    "            if n.left is None:\n",
    "                node = Node('red', None, None, n, value=val)\n",
    "                set_left(node=node, parent=n)\n",
    "                return node\n",
    "            else:\n",
    "                n = n.left\n",
    "        else:\n",
    "            if n.right is None:\n",
    "                node = Node('red', None, None, n, value=val)\n",
    "                set_right(node=node, parent=n)\n",
    "                return node\n",
    "            else:\n",
    "                n = n.right\n",
    "    \n",
    "    \n",
    "def insert(T, v):\n",
    "    # Perform a binary insert.\n",
    "    x = binary_insert(T, v)\n",
    "    \n",
    "    \n",
    "    # Now correct violations of the red-black property.\n",
    "    while x != T.root and x.parent.color == 'red':\n",
    "        if x.parent == x.parent.parent.left:\n",
    "            y = x.parent.parent.right\n",
    "            \n",
    "            # Case 1, both the node and its parent are red. \n",
    "            # Simply recolor.\n",
    "            # This can move the issue further up the tree, which is why we have a while loop.\n",
    "            # So we can fix those issues as well.\n",
    "            if y.color == 'red':\n",
    "                x.parent.color = 'black'\n",
    "                y.color = 'black'\n",
    "                x.parent.parent.color = 'red'\n",
    "                x = x.parent.parent\n",
    "                \n",
    "            else:\n",
    "                # Case 2:\n",
    "                # * x is not the tree root.\n",
    "                # * x's parent is red.\n",
    "                # * x's parent is to the left of its grandparent.\n",
    "                # * x's uncle (its parent's right sibling) is black.\n",
    "                # * x is on its parent's right.\n",
    "                # \n",
    "                # x's parent is red but its uncle is black, which violates invariant four.\n",
    "                # We can fix this by performing a left rotation.                \n",
    "                if x == x.parent.right:\n",
    "                    x = x.parent\n",
    "                    left_rotate(T, x)\n",
    "                \n",
    "                # Case 3. After fixing case 2, we will always end up with a doubled red node further up the tree.\n",
    "                # Perform a right rotation to fix this.\n",
    "                x.parent.color = 'black'\n",
    "                x.parent.parent.color = 'red'\n",
    "                right_rotate(T, x.parent.parent)\n",
    "                \n",
    "        else:  # x.parent == x.parent.parent.right\n",
    "            # Same logic, just with right and left exchanged.\n",
    "            y = x.parent.parent.left\n",
    "            \n",
    "            if y.color == 'red':\n",
    "                x.parent.color = 'black'\n",
    "                y.color = 'black'\n",
    "                x.parent.parent.color = 'red'\n",
    "                x = x.parent.parent\n",
    "                \n",
    "            else:\n",
    "                if x == x.parent.left:\n",
    "                    x = x.parent\n",
    "                    right_rotate(T, x)\n",
    "                    \n",
    "                x.parent.color = 'black'\n",
    "                x.parent.parent.color = 'red'\n",
    "                left_rotate(T, x.parent.parent)\n",
    "    \n",
    "    T.root.color = 'black'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test suite for the swap operation.\n",
    "\n",
    "def test_set_left_clean():\n",
    "    r = Node('red', None, None, None)\n",
    "    t = Tree(r)\n",
    "    n = Node('red', None, None, None)\n",
    "    set_left(n, r)\n",
    "    assert r.left == n\n",
    "    assert n.parent == r\n",
    "    assert r.right == None\n",
    "    assert n.left == None\n",
    "    assert n.right == None\n",
    "    \n",
    "    \n",
    "def test_set_right_clean():\n",
    "    r = Node('black', None, None, None)\n",
    "    t = Tree(r)\n",
    "    n = Node('red', None, None, None)\n",
    "    set_right(n, r)\n",
    "    assert r.right == n\n",
    "    assert n.parent == r\n",
    "    assert r.left == None\n",
    "    assert n.left == None\n",
    "    assert n.right == None\n",
    "    \n",
    "    \n",
    "def test_set_left_overwrite():\n",
    "    r = Node('black', None, None, None)\n",
    "    t = Tree(r)\n",
    "    n1 = Node('red', None, None, None)\n",
    "    n2 = Node('red', None, None, None)\n",
    "    \n",
    "    r.left = n1\n",
    "    n1.parent = r\n",
    "    \n",
    "    set_left(n2, r)\n",
    "    assert r.left == n2\n",
    "    assert n2.parent == r\n",
    "    assert n1.parent == None\n",
    "\n",
    "    \n",
    "def test_set_right_overwrite():\n",
    "    r = Node('black', None, None, None)\n",
    "    t = Tree(r)\n",
    "    n1 = Node('red', None, None, None)\n",
    "    n2 = Node('red', None, None, None)\n",
    "    \n",
    "    r.right = n1\n",
    "    n1.parent = r\n",
    "    \n",
    "    set_right(n2, r)\n",
    "    assert r.right == n2\n",
    "    assert n2.parent == r\n",
    "    assert n1.parent == None\n",
    "\n",
    "\n",
    "def test_set_left_swap():\n",
    "    r = Node('black', None, None, None)\n",
    "    t = Tree(r)\n",
    "    n = Node('red', None, None, None)\n",
    "    \n",
    "    r.right = n\n",
    "    n.parent = r\n",
    "    \n",
    "    set_left(n, r)\n",
    "    assert r.left == n\n",
    "    assert n.parent == r\n",
    "    assert r.right == None\n",
    "    \n",
    "    \n",
    "def test_set_right_swap():\n",
    "    r = Node('black', None, None, None)\n",
    "    t = Tree(r)\n",
    "    n = Node('red', None, None, None)\n",
    "    \n",
    "    r.left = n\n",
    "    n.parent = r\n",
    "    \n",
    "    set_right(n, r)\n",
    "    assert r.right == n\n",
    "    assert n.parent == r\n",
    "    assert r.left == None\n",
    "\n",
    "    \n",
    "    \n",
    "test_set_left_clean()\n",
    "test_set_right_clean()\n",
    "test_set_left_overwrite()\n",
    "test_set_right_overwrite()\n",
    "test_set_left_swap()\n",
    "test_set_right_swap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_left_rotate_nonroot():\n",
    "    r = Node('black', None, None, None)\n",
    "    T = Tree(r)\n",
    "    x = Node('red', None, None, None)\n",
    "    y = Node('black', None, None, None)\n",
    "    \n",
    "    r.right = x\n",
    "    x.parent = r\n",
    "    x.right = y\n",
    "    y.parent = x\n",
    "\n",
    "    left_rotate(T, x)\n",
    "    \n",
    "    assert T.root == r\n",
    "    assert T.root.right == y\n",
    "    assert y.left == x\n",
    "    \n",
    "    \n",
    "def test_left_rotate_root():\n",
    "    r = x = Node('black', None, None, None)\n",
    "    T = Tree(r)\n",
    "    y = Node('black', None, None, None)\n",
    "    \n",
    "    r.right = y\n",
    "    y.parent = r\n",
    "    \n",
    "    left_rotate(T, r)\n",
    "    \n",
    "    assert T.root == y\n",
    "    assert T.root.left == r\n",
    "    \n",
    "    \n",
    "def test_right_rotate_nonroot():\n",
    "    r = Node('black', None, None, None)\n",
    "    T = Tree(r)\n",
    "    x = Node('red', None, None, None)\n",
    "    y = Node('black', None, None, None)\n",
    "    \n",
    "    r.left = x\n",
    "    x.parent = r\n",
    "    x.left = y\n",
    "    y.parent = x\n",
    "\n",
    "    right_rotate(T, x)\n",
    "    \n",
    "    assert T.root == r\n",
    "    assert T.root.left == y\n",
    "    assert y.right == x\n",
    "    \n",
    "    \n",
    "def test_right_rotate_root():\n",
    "    r = x = Node('black', None, None, None)\n",
    "    T = Tree(r)\n",
    "    y = Node('black', None, None, None)\n",
    "    \n",
    "    r.left = y\n",
    "    y.parent = r\n",
    "    \n",
    "    right_rotate(T, r)\n",
    "    \n",
    "    assert T.root == y\n",
    "    assert T.root.right == r\n",
    "    \n",
    "    \n",
    "test_left_rotate_nonroot()\n",
    "test_left_rotate_root()\n",
    "test_right_rotate_nonroot()\n",
    "test_right_rotate_root()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_binary_insert_left():\n",
    "    r = Node('black', None, None, None, value=1)\n",
    "    T = Tree(r)\n",
    "    binary_insert(T, 0)\n",
    "    \n",
    "    assert T.root.left.value == 0\n",
    "    \n",
    "    \n",
    "def test_binary_insert_right():\n",
    "    r = Node('black', None, None, None, value=1)\n",
    "    T = Tree(r)\n",
    "    binary_insert(T, 2)\n",
    "    \n",
    "    assert T.root.right.value == 2\n",
    "\n",
    "    \n",
    "def test_binary_insert_dig():\n",
    "    r = Node('black', None, None, None, value=1)\n",
    "    T = Tree(r)\n",
    "    binary_insert(T, 0)\n",
    "    binary_insert(T, 2)\n",
    "    \n",
    "    binary_insert(T, 3)\n",
    "    binary_insert(T, -1)\n",
    "    \n",
    "    assert T.root.right.right.value == 3\n",
    "    assert T.root.left.left.value == -1\n",
    "\n",
    "    \n",
    "test_binary_insert_left()\n",
    "test_binary_insert_right()\n",
    "test_binary_insert_dig()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def test_insert():\n",
    "    rightright = Node('red', None, None, None, value=15)\n",
    "    right = Node('black', None, rightright, None, value=14)\n",
    "    rightright.parent = right\n",
    "    \n",
    "    leftleft = Node('black', None, None, None, value=1)\n",
    "    leftrightright = Node('red', None, None, None, value=8)\n",
    "    \n",
    "    # leftrightleftleft = Node('red', None, None, value=4)\n",
    "    leftrightleft = Node('red', None, None, None, value=5)\n",
    "    leftright = Node('black', leftrightleft, leftrightright, None, value=7)\n",
    "    leftrightleft.parent = leftright\n",
    "    leftrightright.parent = leftright\n",
    "    \n",
    "    left = Node('red', leftleft, leftright, None, value=2)\n",
    "    leftleft.parent = left\n",
    "    leftright.parent = left\n",
    "    \n",
    "    root = Node('black', left, right, None, value=11)\n",
    "    left.parent = root\n",
    "    right.parent = root\n",
    "    \n",
    "    t = Tree(root)\n",
    "    \n",
    "    insert(t, 4)\n",
    "    \n",
    "test_insert()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* This case above emulates the case given in the reference.\n",
    "* Not included: deletion, which is also $O(\\log{n})$ and \"only slightly more complicated than insertion\"."
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}