{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def z(s):\n",
    "    ''' Use Z-algorithm to preprocess given string.  See\n",
    "        Gusfield for complete description of algorithm. '''\n",
    "    \n",
    "    Z = [len(s)] + [0] * len(s)\n",
    "    assert len(s) > 1\n",
    "    \n",
    "    # Initial comparison of s[1:] with prefix\n",
    "    for i in range(1, len(s)):\n",
    "        if s[i] == s[i-1]:\n",
    "            Z[1] += 1\n",
    "        else:\n",
    "            break\n",
    "    \n",
    "    r, l = 0, 0\n",
    "    if Z[1] > 0:\n",
    "        r, l = Z[1], 1\n",
    "    \n",
    "    for k in range(2, len(s)):\n",
    "        assert Z[k] == 0\n",
    "        if k > r:\n",
    "            # Case 1\n",
    "            for i in range(k, len(s)):\n",
    "                if s[i] == s[i-k]:\n",
    "                    Z[k] += 1\n",
    "                else:\n",
    "                    break\n",
    "            r, l = k + Z[k] - 1, k\n",
    "        else:\n",
    "            # Case 2\n",
    "            # Calculate length of beta\n",
    "            nbeta = r - k + 1\n",
    "            Zkp = Z[k - l]\n",
    "            if nbeta > Zkp:\n",
    "                # Case 2a: Zkp wins\n",
    "                Z[k] = Zkp\n",
    "            else:\n",
    "                # Case 2b: Compare characters just past r\n",
    "                nmatch = 0\n",
    "                for i in range(r+1, len(s)):\n",
    "                    if s[i] == s[i - k]:\n",
    "                        nmatch += 1\n",
    "                    else:\n",
    "                        break\n",
    "                l, r = k, r + nmatch\n",
    "                Z[k] = r - k + 1\n",
    "    return Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[11, 0, 0, 1, 0, 1, 0, 4, 0, 0, 1, 0]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z('abracadabra')\n",
    "#  abracadabra (11)\n",
    "#     a        (1)\n",
    "#       a      (1)\n",
    "#         abra (4)\n",
    "#            a (1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[5, 4, 3, 2, 1, 0]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z('aaaaa')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zMatch(p, t):\n",
    "    s = p + \"$\" + t\n",
    "    Z = z(s)\n",
    "    occurrences = []\n",
    "    for i in range(len(p) + 1, len(s)):\n",
    "        if Z[i] >= len(p):\n",
    "            occurrences.append(i - (len(p) + 1))\n",
    "    return occurrences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[9]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zMatch('needle', 'haystack needle haystack')\n",
    "#                 012345678901234567890123\n",
    "#                           1         2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[9, 16]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zMatch('needle', 'haystack needle needle')\n",
    "#                 0123456789012345678901\n",
    "#                           1         2"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}