{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import re\n",
    "import itertools\n",
    "from   collections import defaultdict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Count Things\n",
    "\n",
    "## Student Records: Late, Absent, Present\n",
    "\n",
    "Consider this problem:\n",
    "\n",
    "> (1) Students at a school must meet with the guidance counselor if they have two absences, or three consecutive late days. Each student's attendance record consists of a string of 'A' for absent, 'L' for late, or 'P' for present. For example: \"LAPLPA\" requires a meeting (because there are two absences), and \"LAPLPL\" is OK (there are three late days, but they are not consecutive). Write a function that takes such a string as input and that true if the student's record is OK. \n",
    "\n",
    "> (2) Write a function to calculate the number of attendance records of length N that are OK.\n",
    "\n",
    "For part (1), the simplest approach is to use `re.search`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ok(record: str) -> bool: return not re.search(r'LLL|A.*A', record)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ok'"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def test_ok():\n",
    "    assert     ok(\"LAPLLP\")\n",
    "    assert not ok(\"LAPLLL\")  # 3 Ls in a row\n",
    "    assert not ok(\"LAPLLA\")  # 2 As overall\n",
    "    assert     ok(\"APLLPLLP\")\n",
    "    assert not ok(\"APLLPLLL\") # 3 Ls in a row\n",
    "    assert not ok(\"APLLPLLA\") # 2 As overall\n",
    "    return 'ok'\n",
    "    \n",
    "test_ok()  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For part (2), I'll start with a simple (but slow) solution that enumerates all possible strings (using `itertools.product`) and checks each one. I use the `quantify` recipe from `itertools` to count them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def total_ok_slow(N: int) -> int:\n",
    "    \"How many strings over 'LAP' of length N are ok?\"\n",
    "    return quantify(all_strings('LAP', N), ok)\n",
    "\n",
    "def quantify(iterable, pred=bool):\n",
    "    \"Count how many times the predicate is true of items in iterable.\"\n",
    "    return sum(map(pred, iterable))\n",
    "\n",
    "cat = ''.join\n",
    "\n",
    "def all_strings(alphabet, N): return map(cat, itertools.product(alphabet, repeat=N))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 1,\n",
       " 1: 3,\n",
       " 2: 8,\n",
       " 3: 19,\n",
       " 4: 43,\n",
       " 5: 94,\n",
       " 6: 200,\n",
       " 7: 418,\n",
       " 8: 861,\n",
       " 9: 1753,\n",
       " 10: 3536}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{N: total_ok_slow(N) for N in range(11)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks good, but\n",
    "I will need a more efficient algorithm to handle large values of *N*. Here's how I think about it:\n",
    "\n",
    "* I can't enumerate all the strings; there are too many of them, 3<sup>N</sup>. \n",
    "* Even if I only enumerate the ok strings, there are still too many, O(2<sup>N</sup>).\n",
    "* Instead, I'll want to keep track of a *summary* of all the ok strings of length *N*, and use that to quickly compute a summary of the ok strings of length *N*+1. I recognize this as a *[dynamic programming](https://en.wikipedia.org/wiki/Dynamic_programming)* approach.\n",
    "\n",
    "* What is in the summary? A list of all ok strings is too much. A count of the number of ok strings is not enough. Instead, I will group together the strings that have the same number of `'A'` characters in them, and the same number of consecutive `'L'` characters at the end of the string, and count them.  I don't need to count strings that have two or more `'A'` characters, or 3 consecutive `'L'` characters anywhere in the string. And I don't need to worry about runs of 1 or 2 `'L'` characters embedded in the middle of the string. So the summary is a mapping of the form `{(A, L): count, ...}`. \n",
    "\n",
    "* For *N* = 2, the summary looks like this:\n",
    "\n",
    "\n",
    "     #(A, L): count\n",
    "     {(0, 0): 2, # LP, PP\n",
    "      (0, 1): 1, # PL\n",
    "      (0, 2): 1, # LL\n",
    "      (1, 0): 1, # AP, LA, PA\n",
    "      (1, 1): 1} # AL\n",
    "\n",
    "\n",
    "Here is a function to create the summary for `N+1`, given the summary for `N`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def next_summary(prev_summary):\n",
    "    \"Given a summary of the form {(A, L): count, ...}, return the summary for strings one character longer.\"\n",
    "    summary = defaultdict(int)\n",
    "    for (A, L), c in prev_summary.items():\n",
    "            if A < 1: summary[A+1, 0] += c # transition with 'A'\n",
    "            if L < 2: summary[A, L+1] += c # transition with 'L'\n",
    "            summary[A, 0] += c             # transition with 'P'\n",
    "    return summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For `N = 0`, the summary is `{(0, 0): 1}`, because there is one string, the empty string, which has no `'A'` nor `'L'`. From there we can proceed in a \"bottom-up\" fashion to compute the total number of OK strings for any value of `N`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a \"bottom-up\" approach for `total_ok` that starts at `0` and works up to `N`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def total_ok(N):\n",
    "    \"How many strings of length N are ok?\"\n",
    "    summary = {(0, 0): 1}\n",
    "    for _ in range(N):\n",
    "        summary = next_summary(summary)\n",
    "    return sum(summary.values()) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this to go way beyond what we could do with `total_ok_slow`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.43 ms, sys: 50 µs, total: 2.48 ms\n",
      "Wall time: 2.48 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "5261545087067582125179062608958232695543100705754634272071166414871321070487675367"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time total_ok(300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are over 10<sup>80</sup> ok strings of length 300; more than the number of atoms in the universe. But it only took around a millisecond to count them.\n",
    "\n",
    "Dynamic programming can be done top-down (where we start at `N` and work down to `0`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def total_ok(N):\n",
    "    \"How many strings of length N are ok?\"\n",
    "    return sum(summary_for(N).values())\n",
    "    \n",
    "def summary_for(N): return ({(0, 0): 1} if N == 0 else next_summary(summary_for(N - 1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.33 ms, sys: 538 µs, total: 2.87 ms\n",
      "Wall time: 4.83 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "5261545087067582125179062608958232695543100705754634272071166414871321070487675367"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time total_ok(300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make sure we're getting the same results as before, and take a look at the summaries for the first few values of `N`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0    1 {(0, 0): 1}\n",
      " 1    3 {(0, 1): 1, (1, 0): 1, (0, 0): 1}\n",
      " 2    8 {(0, 1): 1, (1, 0): 3, (0, 0): 2, (0, 2): 1, (1, 1): 1}\n",
      " 3   19 {(0, 1): 2, (1, 2): 1, (0, 0): 4, (1, 0): 8, (0, 2): 1, (1, 1): 3}\n",
      " 4   43 {(0, 1): 4, (1, 2): 3, (0, 0): 7, (1, 0): 19, (0, 2): 2, (1, 1): 8}\n",
      " 5   94 {(0, 1): 7, (1, 2): 8, (0, 0): 13, (1, 0): 43, (0, 2): 4, (1, 1): 19}\n",
      " 6  200 {(0, 1): 13, (1, 2): 19, (0, 0): 24, (1, 0): 94, (0, 2): 7, (1, 1): 43}\n",
      " 7  418 {(0, 1): 24, (1, 2): 43, (0, 0): 44, (1, 0): 200, (0, 2): 13, (1, 1): 94}\n",
      " 8  861 {(0, 1): 44, (1, 2): 94, (0, 0): 81, (1, 0): 418, (0, 2): 24, (1, 1): 200}\n",
      " 9 1753 {(0, 1): 81, (1, 2): 200, (0, 0): 149, (1, 0): 861, (0, 2): 44, (1, 1): 418}\n",
      "10 3536 {(0, 1): 149, (1, 2): 418, (0, 0): 274, (1, 0): 1753, (0, 2): 81, (1, 1): 861}\n"
     ]
    }
   ],
   "source": [
    "for N in range(11): \n",
    "    assert total_ok(N) == total_ok_slow(N)\n",
    "    print('{:2} {:4} {}'.format(N, total_ok(N), dict(summary_for(N))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# >>> Count Strings with Alphabetic First Occurences\n",
    "\n",
    "> Given an alphabet of length k, how many strings of length k can be formed such that the first occurrences of each character in the string are a prefix of the alphabet?\n",
    "\n",
    "Let's first make sure we understand the problem. I will choose to represent a string as a list of integers, like `[0, 1, 2]` rather than as a `str` like `\"abc\"`, and the alphabet will always be `range(k)`. So, the string `[0, 1, 0, 2]` would be valid, because the first occurrences are `[0, 1, 2]`, but `[0, 1, 0, 3]` would not be valid, since `[0, 1, 3]` is not a prefix of `range(4)`. I'll define three key concepts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def valid(s): return is_prefix(first_occurrences(s))\n",
    "\n",
    "def is_prefix(s): return s == list(range(len(s)))\n",
    "\n",
    "def first_occurrences(s):\n",
    "    \"The unique elements of s, in the order they appear.\" \n",
    "    firsts = []\n",
    "    for x in s:\n",
    "        if x not in firsts: firsts.append(x)\n",
    "    return firsts "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ok'"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def test(): #    s             firsts(s)     valid(s)\n",
    "    assert test1([0, 1, 2],    [0, 1, 2],    True)  \n",
    "    assert test1([0, 0, 0],    [0],          True)      \n",
    "    assert test1([1],          [1],          False)      \n",
    "    assert test1([0, 1, 3],    [0, 1, 3],    False)\n",
    "    assert test1([0, 1, 3, 2], [0, 1, 3, 2], False)\n",
    "    assert test1([0, 1, 0, 1, 0, 2, 1], [0, 1, 2], True)\n",
    "    assert test1([0, 1, 0, 2, 1, 3, 1, 2, 5, 4, 3], [0, 1, 2, 3, 5, 4], False)\n",
    "    return 'ok'\n",
    "\n",
    "def test1(s, firsts, is_valid):\n",
    "    \"\"\"Test whether first_occurrences(s) == firsts and valid(s) == is_valid\"\"\"\n",
    "    return first_occurrences(s) == firsts and valid(s) == is_valid\n",
    "    \n",
    "test()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, I will solve the problem in a slow but sure way: generate all possible strings, then count the number that are valid. The complexity of this algorithm is $O(k^{k+1})$, because there are $k^k$ strings, and to validate a string requires looking at all $k$ characters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 1, 2, 5, 15, 52, 203]"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import itertools \n",
    "\n",
    "all_strings = itertools.product\n",
    "\n",
    "def how_many_slow(k): \n",
    "    \"\"\"Count the number of valid strings. (Try all possible strings.)\"\"\"\n",
    "    return sum(valid(s) for s in all_strings(range(k), repeat=k))\n",
    "\n",
    "[how_many_slow(k) for k in range(7)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's think about how to speed that up.I don't want to have to consider every possible string, because there are too many ($k^k$) of them. Can I group together many strings and just count the number of them, without enumerating each one? For example, if I knew there were 52 valid strings of length $k-1$ (and didn't know anything else about them), can I tell how many valid strings of length $k$ there are? I don't see a way to do this, because the number of ways to extend a valid string is dependent on the number of distinct characters in the string. If a string has $m$ distinct characters, then I can extend it by repeating any of those $m$ characters, or by introducing a first occurrence of character number $m+1$.\n",
    "\n",
    "So I need to keep track of the number of valid strings of length $k$ that have exactly $m$ distinct characters (which, by definition, must be `range(m)`). I'll call that number `C(k, m)`. Then I can define `how_many(k)` as the sum over all values of `m` of `C(k, m)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from functools import lru_cache\n",
    "    \n",
    "@lru_cache()\n",
    "def C(k, m):\n",
    "    \"Count the number of valid strings of length k, that use m distinct characters.\"\n",
    "    return (1 if k == 0 == m else\n",
    "            0 if k == 0 != m else\n",
    "            C(k-1, m) * m + C(k-1, m-1)) # m ways to add an old character; 1 way to add new\n",
    "\n",
    "def how_many(k): return sum(C(k, m) for m in range(k+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "47585391276764833658790768841387207826363669686825611466616334637559114497892442622672724044217756306953557882560751"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "how_many(100)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "assert(all(how_many(k) == how_many_slow(k) for k in range(7)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0 10^0   1\n",
      "  1 10^0   1\n",
      "  2 10^0   2\n",
      "  3 10^1   5\n",
      "  4 10^1   15\n",
      "  5 10^2   52\n",
      "  6 10^2   203\n",
      "  7 10^3   877\n",
      "  8 10^4   4140\n",
      "  9 10^4   21147\n",
      " 10 10^5   115975\n",
      " 20 10^14  51724158235372\n",
      " 30 10^24  846749014511809332450147\n",
      " 40 10^35  157450588391204931289324344702531067\n",
      " 50 10^47  185724268771078270438257767181908917499221852770\n",
      " 60 10^60  976939307467007552986994066961675455550246347757474482558637\n",
      " 70 10^73  18075003898340511237556784424498369141305841234468097908227993035088029195\n",
      " 80 10^87  991267988808424794443839434655920239360814764000951599022939879419136287216681744888844\n",
      " 90 10^101 141580318123392930464192819123202606981284563291786545804370223525364095085412667328027643050802912567\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "\n",
    "for k in itertools.chain(range(10), range(10, 100, 10)):\n",
    "    n = how_many(k)\n",
    "    print('{:3} 10^{:<3} {:d}'.format(k, round(math.log10(n)), n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-11 -9 5\n",
      "-5 9 11\n",
      "-1 4 11\n",
      "done 20\n"
     ]
    }
   ],
   "source": [
    "N = 20\n",
    "for a, b, p in itertools.combinations(range(-N, N), 3):\n",
    "    if (b + p) * (a + p) * (a + b) == 0: continue\n",
    "    if a/(b + p) + b/(a + p) + p/(a + b) == 4:\n",
    "        print(a, b, p)\n",
    "print('done', N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "(a^3 + a^2 b + a^2 p + a b^2 + 3 a b p + a p^2 + b^3 + b^2 p + b p^2 + p^3)/((a + b) (a + p) (b + p)) = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 25 s, sys: 106 ms, total: 25.1 s\n",
      "Wall time: 25.3 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "945"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sides = range(1, 11)\n",
    "from itertools import permutations\n",
    "\n",
    "def sortuple(items): return tuple(sorted(items))\n",
    "\n",
    "def sets_of_rectangles(sides): return set(*map(set_of_rectangles, permutations(sides)))\n",
    "\n",
    "def set_of_rectangles(sides): return sortuple(sortuple(sides[i:i+2]) for i in range(0, len(sides), 2))\n",
    "\n",
    "%time len(sets_of_rectangles(sides))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((1, 2), (3, 4), (5, 6), (7, 8), (9, 10))"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from itertools import zip_longest\n",
    "\n",
    "def grouper(iterable, n, fillvalue=None):\n",
    "    \"Collect data into fixed-length chunks or blocks\"\n",
    "    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx\"\n",
    "    args = [iter(iterable)] * n\n",
    "    return zip_longest(*args, fillvalue=fillvalue)\n",
    "\n",
    "def grouper(iterable, n):\n",
    "    \n",
    "\n",
    "tuple(grouper(sides, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{frozenset({(1, 3), (2, 5), (4, 6)}),\n",
       " frozenset({(1, 2), (3, 5), (4, 6)}),\n",
       " frozenset({(1, 3), (2, 6), (4, 5)}),\n",
       " frozenset({(1, 6), (2, 3), (4, 5)}),\n",
       " frozenset({(1, 4), (2, 3), (5, 6)}),\n",
       " frozenset({(1, 6), (2, 4), (3, 5)}),\n",
       " frozenset({(1, 6), (2, 5), (3, 4)}),\n",
       " frozenset({(1, 5), (2, 3), (4, 6)}),\n",
       " frozenset({(1, 5), (2, 6), (3, 4)}),\n",
       " frozenset({(1, 3), (2, 4), (5, 6)}),\n",
       " frozenset({(1, 4), (2, 6), (3, 5)}),\n",
       " frozenset({(1, 2), (3, 4), (5, 6)}),\n",
       " frozenset({(1, 5), (2, 4), (3, 6)}),\n",
       " frozenset({(1, 2), (3, 6), (4, 5)}),\n",
       " frozenset({(1, 4), (2, 5), (3, 6)})}"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sets_of_rectangles(sides)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
