{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "dfda8eaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Suffix tree built with simple O(m^2)-time algorithm,\n",
    "# adding a suffix at a time \n",
    "class SuffixTree(object):\n",
    "    \n",
    "  class Node(object):\n",
    "    # class variable counting the number created of nodes\n",
    "    node_count = 0\n",
    "    verbose = False \n",
    "    def __init__(self, depth, off, ln):\n",
    "      self.depth = depth # string depth \n",
    "      self.off = off # offset into T of substring on incoming edge\n",
    "      self.ln = ln   # length of substring on incoming edge\n",
    "      self.out = {}  # outgoing edges; maps characters to nodes\n",
    "      self.id = SuffixTree.Node.node_count # give unique id to each node\n",
    "      SuffixTree.Node.node_count += 1\n",
    "    def __repr__(self):\n",
    "      return f\"string_depth:{self.depth:2} out:{sorted(self.out.keys())}, edge(off:{self.off} len:{self.ln})\"\n",
    "  \n",
    "  # create the suffix tree for the string t, add $ if necessary\n",
    "  def __init__(self, t):\n",
    "    \"\"\" Make suffix tree (without suffix links) from t in quadratic time\n",
    "        and linear space adding one suffix at a time\"\"\"\n",
    "    if t[-1]!='$':    \n",
    "      t += '$'     # add unique EOS symbol \n",
    "    self.t = t     # save string in istance variable \n",
    "    self.root = self.Node(0, 0, 0)               # create root\n",
    "    self.root.out[t[0]] = self.Node(len(t), 0, len(t)) # create leaf for while t\n",
    "    for i in range(1, len(t)):\n",
    "      cur = self.root\n",
    "      j = i               # add t[i:] to suffix tree\n",
    "      while j < len(t):\n",
    "        if t[j] in cur.out:\n",
    "          child = cur.out[t[j]]\n",
    "          # compute child label\n",
    "          lab = t[child.off:child.off+child.ln]\n",
    "          k = j+1 # Walk along edge\n",
    "          while k-j < len(lab) and t[k] == lab[k-j]:\n",
    "              k += 1\n",
    "          if k-j == len(lab):\n",
    "              cur = child # exhausted the edge\n",
    "              j = k\n",
    "          else:\n",
    "              # fell off in middle of edge\n",
    "              cExist, cNew = lab[k-j], t[k]\n",
    "              mid = self.Node(cur.depth + k-j, child.off, k-j)\n",
    "              mid.out[cNew] = self.Node(mid.depth + len(t[k:]), k, len(t[k:]))\n",
    "              mid.out[cExist] = child\n",
    "              child.off += (k-j)\n",
    "              child.ln -= (k-j)\n",
    "              cur.out[t[j]] = mid\n",
    "        else:\n",
    "          # Create a new edge hanging off of this node\n",
    "          cur.out[t[j]] = self.Node(cur.depth + len(t[j:]), j, len(t[j:]))\n",
    "\n",
    "  def followPath(self, s):\n",
    "    \"\"\" Follow path given by s.  If we fall off tree, return None.  If we\n",
    "        finish mid-edge, return (node, offset) where 'node' is child and\n",
    "        'offset' is label offset.  If we finish on a node, return (node,\n",
    "        None). \"\"\"\n",
    "    cur = self.root\n",
    "    i = 0    # i index inside s (the pattern)\n",
    "    while i < len(s):\n",
    "      c = s[i]\n",
    "      if c not in cur.out:\n",
    "          return (None, None) # fell off at a node\n",
    "      child = cur.out[s[i]]\n",
    "      lab = self.t[child.off:child.off+child.ln]\n",
    "      j = i+1\n",
    "      while j-i < len(lab) and j < len(s) and s[j] == lab[j-i]:\n",
    "          j += 1\n",
    "      if j-i == len(lab):\n",
    "          cur = child # exhausted edge\n",
    "          i = j\n",
    "      elif j == len(s): \n",
    "          return (child, j-i) # exhausted query string in middle of edge\n",
    "      else:\n",
    "          return (None, None) # fell off in the middle of the edge\n",
    "    return (cur, None) # exhausted query string at internal node    \n",
    "  \n",
    "  def report_occ(self,s):\n",
    "    \"\"\"Report the positions of the occurrences of s in t\n",
    "       note occurrences are reported in lexicographic order\"\"\"\n",
    "    n,o = self.followPath(s)\n",
    "    if n == None:\n",
    "      return []\n",
    "    occ = []\n",
    "    def __visit(n):\n",
    "      if len(n.out) == 0:\n",
    "        # leaf node, append suffix id to occ\n",
    "        occ.append(len(self.t) - n.depth)\n",
    "      # visit children in lexicographical order\n",
    "      for c, child in sorted(n.out.items()):\n",
    "        __visit(child)\n",
    "    __visit(n)\n",
    "    return occ\n",
    "    \n",
    "  def __repr__(self):\n",
    "    \"\"\" return a string containing the result of \n",
    "        a pre-order visit. Left indentation shows\n",
    "        the tree depth of a node\"\"\"\n",
    "    return self.repr_ric(self.root,0)\n",
    "\n",
    "  def print_leaves(self):\n",
    "    \"\"\" print all leaves of self\"\"\"\n",
    "    self.print_leaves_ric(self.root)\n",
    "\n",
    "  def str_node(self,n):\n",
    "    \"\"\"return the string representation of a node n of self\n",
    "       the label is the string in the root to n path, with a |\n",
    "       marking the begining of the content of the edge leading to n\"\"\"\n",
    "    return f\"depth:{n.depth} out:{str(sorted(n.out.keys()))} \"+ \\\n",
    "           f\"label:{self.t[-n.depth:-n.ln]}|{self.t[n.off:n.off+n.ln:]}\"\n",
    "\n",
    "  # ---- auxiliary recursive functions ----    \n",
    "  def repr_ric(self,n,td):\n",
    "    out = td*\" \" + self.str_node(n)+\"\\n\"\n",
    "    for c in sorted(n.out):\n",
    "      out += self.repr_ric(n.out[c],td+1)\n",
    "    return out\n",
    "  \n",
    "  def print_leaves_ric(self,n):\n",
    "    if len(n.out) == 0:\n",
    "      print(self.str_node(n))\n",
    "    else:\n",
    "      for c in sorted(n.out):\n",
    "        self.print_leaves_ric(n.out[c])\n",
    "  \n",
    "  def saLcp(self):\n",
    "    # Return suffix array and LCP array corresponding to this\n",
    "    # suffix tree.  self.root is root, self.t is the text.\n",
    "    sa, lcp1 = [], [-1]\n",
    "    # visit suffix tree in lexicographic order\n",
    "    def __visit(n):\n",
    "      if len(n.out) == 0:\n",
    "        # leaf node: record offset and LCP with previous leaf\n",
    "        sa.append(len(self.t) - n.depth)\n",
    "      else:\n",
    "        # visit children in lexicographical order\n",
    "        first = True\n",
    "        for c, child in sorted(n.out.items()):\n",
    "          if first: first = False\n",
    "          else: lcp1.append(n.depth) \n",
    "          __visit(child)\n",
    "    # end of __visit  \n",
    "    __visit(self.root)\n",
    "    return sa, lcp1\n",
    "\n",
    "  def bwt(self):\n",
    "    # return BWT for self.t doing a lexicographic visit\n",
    "    _bwt = []\n",
    "    # visit suffix tree in lexicographic order\n",
    "    def __visit(n): \n",
    "      if len(n.out) == 0:\n",
    "        # leaf node: add associate symbol to bwt\n",
    "        _bwt.append(self.t[len(self.t)-n.depth-1][0]) \n",
    "      else:\n",
    "        # visit children in lexicographical order\n",
    "        for c, child in sorted(n.out.items()):\n",
    "          __visit(child)\n",
    "    # end of __visit  \n",
    "    __visit(self.root)\n",
    "    return \"\".join(_bwt)\n",
    "\n",
    "# show SA/LCP and list of sorted suffixes\n",
    "def show_saLcp(t,sa,lcp):\n",
    "  if t[-1]!='$': \n",
    "    t+='$'   #  add final dollar if needed\n",
    "  assert len(t) == len(sa)\n",
    "  for i in range(len(sa)):\n",
    "    print(f\"{sa[i]:2} {lcp[i]:2} {t[sa[i]:]}\")\n",
    "  \n",
    "# simple construction of the BWT from the SA  \n",
    "def bwt(t,sa):\n",
    "  assert len(t) == len(sa)\n",
    "  bwt = \"\".join([t[sa[i]-1] for i in range(len(sa))])\n",
    "  return bwt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4e95287b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "string_depth: 0 out:['$', 'i', 'm', 'p', 's'], edge(off:0 len:0)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example: suffix tree for mississippi\n",
    "t = \"mississippi$\"\n",
    "st = SuffixTree(t)\n",
    "# root with 5 outgoing edges\n",
    "st.root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "89c933a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "string_depth: 1 out:['i', 's'], edge(off:2 len:1)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# follow edge with label 's' from the root\n",
    "snode = st.root.out[\"s\"]\n",
    "snode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5f2b86f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "string_depth: 3 out:['p', 's'], edge(off:3 len:2)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# follow again edge with label s\n",
    "# note there are 2 symbols on that edge\n",
    "snode.out[\"s\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "cb4a9ffb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- SA and LCP\n",
      "11 -1 $\n",
      "10  0 i$\n",
      " 7  1 ippi$\n",
      " 4  1 issippi$\n",
      " 1  4 ississippi$\n",
      " 0  0 mississippi$\n",
      " 9  0 pi$\n",
      " 8  1 ppi$\n",
      " 6  0 sippi$\n",
      " 3  2 sissippi$\n",
      " 5  1 ssippi$\n",
      " 2  3 ssissippi$\n"
     ]
    }
   ],
   "source": [
    "# build SA and LCP array from Suffix Tree\n",
    "sa, lcp = st.saLcp()\n",
    "# show SA/LCP arrays \n",
    "print(\"--- SA and LCP\")\n",
    "show_saLcp(st.t,sa,lcp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "000d92c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- suffix tree leaves in lexicographic order\n",
      "depth:1 out:[] label:|$\n",
      "depth:8 out:[] label:|+ananas$\n",
      "depth:9 out:[] label:a|+ananas$\n",
      "depth:11 out:[] label:ana|+ananas$\n",
      "depth:13 out:[] label:anana|+ananas$\n",
      "depth:7 out:[] label:anana|s$\n",
      "depth:5 out:[] label:ana|s$\n",
      "depth:3 out:[] label:a|s$\n",
      "depth:14 out:[] label:|banana+ananas$\n",
      "depth:10 out:[] label:na|+ananas$\n",
      "depth:12 out:[] label:nana|+ananas$\n",
      "depth:6 out:[] label:nana|s$\n",
      "depth:4 out:[] label:na|s$\n",
      "depth:2 out:[] label:|s$\n"
     ]
    }
   ],
   "source": [
    "# More complex example\n",
    "t = \"banana+ananas$\"\n",
    "# build suffix tree\n",
    "st =  SuffixTree(t)\n",
    "print(\"--- suffix tree leaves in lexicographic order\");\n",
    "st.print_leaves()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5246f6b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- starting positions of the occurrences of 'ana' from ST\n",
      "[3, 1, 7, 9]\n"
     ]
    }
   ],
   "source": [
    "print(\"--- starting positions of the occurrences of 'ana' from ST\")\n",
    "print(st.report_occ(\"ana\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f5f30f42",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- SA and LCP\n",
      "13 -1 $\n",
      " 6  0 +ananas$\n",
      " 5  0 a+ananas$\n",
      " 3  1 ana+ananas$\n",
      " 1  3 anana+ananas$\n",
      " 7  5 ananas$\n",
      " 9  3 anas$\n",
      "11  1 as$\n",
      " 0  0 banana+ananas$\n",
      " 4  0 na+ananas$\n",
      " 2  2 nana+ananas$\n",
      " 8  4 nanas$\n",
      "10  2 nas$\n",
      "12  0 s$\n"
     ]
    }
   ],
   "source": [
    "# build SA and LCP array from Suffix Tree\n",
    "sa, lcp = st.saLcp()\n",
    "# show SA/LCP arrays \n",
    "print(\"--- SA and LCP\")\n",
    "show_saLcp(st.t,sa,lcp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f648bdf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try reasoning on how to find the longest substring \n",
    "# appearing in both strings using SA/LCP array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "37475761",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- SA and LCP\n",
      "14 -1 $\n",
      " 6  0 +arancia$\n",
      "13  0 a$\n",
      " 5  1 a+arancia$\n",
      " 3  1 ana+arancia$\n",
      " 1  3 anana+arancia$\n",
      " 9  2 ancia$\n",
      " 7  1 arancia$\n",
      " 0  0 banana+arancia$\n",
      "11  0 cia$\n",
      "12  0 ia$\n",
      " 4  0 na+arancia$\n",
      " 2  2 nana+arancia$\n",
      "10  1 ncia$\n",
      " 8  0 rancia$\n"
     ]
    }
   ],
   "source": [
    "t = \"banana+arancia\"\n",
    "# build suffix tree\n",
    "st =  SuffixTree(t)\n",
    "# build SA and LCP array from Suffix Tree\n",
    "sa, lcp = st.saLcp()\n",
    "# show SA/LCP arrays \n",
    "print(\"--- SA and LCP\")\n",
    "show_saLcp(st.t,sa,lcp)"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}