find_ancestor.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from networkx.generators.classic import balanced_tree\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = balanced_tree(r=2, h=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EdgeView([(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6), (3, 7), (3, 8), (4, 9), (4, 10), (5, 11), (5, 12), (6, 13), (6, 14)])"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g.edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NodeView((0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14))"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g.nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 149,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nx.tree.is_tree(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [1, 2]\n",
      "1 [0, 3, 4]\n",
      "2 [0, 5, 6]\n",
      "3 [1, 7, 8]\n",
      "4 [1, 9, 10]\n",
      "5 [2, 11, 12]\n",
      "6 [2, 13, 14]\n",
      "7 [3]\n",
      "8 [3]\n",
      "9 [4]\n",
      "10 [4]\n",
      "11 [5]\n",
      "12 [5]\n",
      "13 [6]\n",
      "14 [6]\n"
     ]
    }
   ],
   "source": [
    "for node in g.nodes:\n",
    "    print(node, list(g.neighbors(node)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code up the tree into a `code` array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "def code_up(g):\n",
    "    \"\"\"Code up a networkx graph representing a tree.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    g : networkx graph\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    coding : np.array\n",
    "        The length of the arrays is the number of vertices in the tree.\n",
    "        The `k`-th element, `coding[k]` is the parent of `k`.\n",
    "        By convention, `coding[0] = -1` which signals the root of the tree.\n",
    "    \"\"\"\n",
    "\n",
    "    coding = np.zeros(g.number_of_nodes(), dtype=int)\n",
    "    coding[0] = -1\n",
    "    for label in range(1, g.number_of_nodes()):\n",
    "        # HACK: use the balanced_tree feature: the parent's\n",
    "        # label is always smaller than that of a child's\n",
    "        coding[label] = min(g.neighbors(label))\n",
    "    return coding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-1,  0,  0,  1,  1,  2,  2,  3,  3,  4,  4,  5,  5,  6,  6])"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "code_up(g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Walk up towards a common ancestor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mark_path(label, coding, visited):\n",
    "    \"\"\"Mark a path from the node `label` up to the root or a nearest marked node.\n",
    "    \n",
    "    The `visited` array is modified in-place.\n",
    "    \"\"\"    \n",
    "    node = label\n",
    "    while True:\n",
    "        if coding[node] == -1:\n",
    "            # this a root\n",
    "            visited[node] = True\n",
    "            return node\n",
    "        \n",
    "        visited[node] = True\n",
    "        node = coding[node]\n",
    "        \n",
    "        if visited[node]:\n",
    "            return node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_ancestor(node_1, node_2, coding):\n",
    "    \"\"\"Find a last common ancestor of `node_1` and `node_2`.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    label_1 : int\n",
    "        The first node\n",
    "    label_2 : int\n",
    "        The second node\n",
    "    coding : array, shape (n,)\n",
    "        The coding of a binary tree of `n` nodes.\n",
    "        The element coding[node] is the label of the parent of `node`.\n",
    "        For a root, `coding[0] = -1`\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    ancestor : int\n",
    "        The label of the common ancestor\n",
    "    visited : array, shape (n,)\n",
    "        visited[k] is True if `k` is a part of the path between the\n",
    "        nodes _1 and _2 and the root.\n",
    "    \"\"\"\n",
    "    visited = np.empty(coding.shape[0], dtype=bool)\n",
    "    visited[:] = False\n",
    "    \n",
    "    _ = mark_path(node_1, coding, visited)\n",
    "    ancestor = mark_path(node_2, coding, visited)\n",
    "    return ancestor, visited"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2,\n",
       " array([ True, False,  True, False, False,  True,  True, False, False,\n",
       "        False, False,  True, False,  True, False]))"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coding = code_up(balanced_tree(r=2, h=3))\n",
    "\n",
    "find_ancestor(11, 13, coding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Move over to cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "\n",
    "# Use %%cython -a to see the annotated C code\n",
    "\n",
    "import numpy as np\n",
    "from numpy cimport npy_int8, npy_intp\n",
    "\n",
    "cimport cython\n",
    "\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "cdef Py_ssize_t mark_path(npy_intp label,\n",
    "                          npy_intp[::1] coding,\n",
    "                          npy_int8[::1] visited):\n",
    "    \"\"\"Mark a path from the node `label` up to the root or a nearest marked node.\n",
    "    \n",
    "    The `visited` array is modified in-place.\n",
    "    \"\"\"\n",
    "    cdef npy_intp node = label\n",
    "    while True:\n",
    "        if coding[node] == -1:\n",
    "            # this a root\n",
    "            visited[node] = True\n",
    "            return node\n",
    "        \n",
    "        visited[node] = True\n",
    "        node = coding[node]\n",
    "        \n",
    "        if visited[node]:\n",
    "            return node\n",
    "\n",
    "\n",
    "def find_ancestor(node_1, node_2, coding):\n",
    "    \"\"\"Find a last common ancestor of `node_1` and `node_2`.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    label_1 : int\n",
    "        The first node\n",
    "    label_2 : int\n",
    "        The second node\n",
    "    coding : array, shape (n,)\n",
    "        The coding of a binary tree of `n` nodes.\n",
    "        The element coding[node] is the label of the parent of `node`.\n",
    "        For a root, `coding[0] = -1`\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    ancestor : int\n",
    "        The label of the common ancestor\n",
    "    visited : array, shape (n,)\n",
    "        visited[k] is True if `k` is a part of the path between the\n",
    "        nodes _1 and _2 and the root.\n",
    "    \"\"\"\n",
    "    visited = np.zeros(coding.shape[0], dtype='int8')\n",
    "    \n",
    "    _ = mark_path(node_1, coding, visited)\n",
    "    ancestor = mark_path(node_2, coding, visited)\n",
    "    return ancestor, visited"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2,\n",
       " array([ True, False,  True, False, False,  True,  True, False, False,\n",
       "        False, False,  True, False,  True, False]))"
      ]
     },
     "execution_count": 158,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coding = code_up(balanced_tree(r=2, h=3))\n",
    "\n",
    "find_ancestor(11, 13, coding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "\n"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}