2021_01_FigurineFiguring.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "burning-shift",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import math\n",
    "import itertools\n",
    "import requests\n",
    "from bs4 import BeautifulSoup\n",
    "import sympy\n",
    "import copy\n",
    "import xlwings as xw\n",
    "from sympy.abc import a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "living-dragon",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Puzzle\n",
      "~~~~~~\n",
      "Jane received 78 figurines as gifts this holiday season:  12 drummers drumming, 11 pipers piping, 10 lords a-leaping, etc., down to 1 partridge in a pear tree.   They are all mixed together in a big bag.  She agrees with her friend Alex that this seems like too many figurines for one person to have, so she decides to give some of her figurines to Alex.   Jane will uniformly randomly pull figurines out of the bag one at a time until she pulls out the partridge in a pear tree, and will give Alex all of the figurines she pulled out of the bag (except the partridge, that’s Jane’s favorite).Ifnis the maximum number of any one type of ornament that Alex gets, what is the expected value ofn, to seven significant figures?\n"
     ]
    }
   ],
   "source": [
    "# I found this stackexchange post very informative\n",
    "# https://math.stackexchange.com/questions/3934750/how-to-solve-this-candy-collectors-puzzle\n",
    "# and the book referred to https://www2.math.upenn.edu/~wilf/DownldGF.html is vey good.\n",
    "\n",
    "print(\"Puzzle\")\n",
    "print(\"~~~~~~\")\n",
    "url='https://www.janestreet.com/puzzles/figurine-figuring/'\n",
    "res = requests.get(url)\n",
    "soup = BeautifulSoup(res.content, 'html.parser')\n",
    "x =[text for text in soup.body.stripped_strings]\n",
    "print(\"\".join(x[7:13]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "widespread-specification",
   "metadata": {},
   "source": [
    "<hr>\n",
    "<b>Generating function is :</b>\n",
    "<br>\n",
    "<br>\n",
    "    $\\Large\\prod_{j=2}^i\\left(\\sum_{k=0}^j \\frac{j!}{(j-k)!} \\frac{a^k}{k!}\\right)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "adjacent-mainstream",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "With a max of 1 of any ball you can draw 11\n",
      "With a max of 2 of any ball you can draw 22\n",
      "With a max of 3 of any ball you can draw 32\n",
      "With a max of 4 of any ball you can draw 41\n",
      "With a max of 5 of any ball you can draw 49\n",
      "With a max of 6 of any ball you can draw 56\n",
      "With a max of 7 of any ball you can draw 62\n",
      "With a max of 8 of any ball you can draw 67\n",
      "With a max of 9 of any ball you can draw 71\n",
      "With a max of 10 of any ball you can draw 74\n",
      "With a max of 11 of any ball you can draw 76\n",
      "With a max of 12 of any ball you can draw 77\n",
      "(12, 78)\n"
     ]
    }
   ],
   "source": [
    "# Generating function using sympy and put the coeffs into an array\n",
    "# remove higher order terms to find probabilites with constraints on \n",
    "# max of any one figure\n",
    "\n",
    "perms=np.zeros((12,78))\n",
    "\n",
    "for cap in range(1,13):\n",
    "    gfa = sympy.poly(1,a)\n",
    "    for j in range(2,13):\n",
    "        gfb = sympy.poly(0,a)\n",
    "        for k in range(0,min(j+1,cap+1)):\n",
    "            gfb += sympy.poly((math.factorial(j)/math.factorial(j-k))*(a**k)/math.factorial(k),a)\n",
    "        gfa *=gfb\n",
    "    print(\"With a max of {} of any ball you can draw {}\".format(cap,len(gfa.monoms())-1))    \n",
    "    perms[cap-1,78-len(gfa.coeffs()):]=gfa.coeffs()\n",
    "\n",
    "print(np.array(perms).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "mineral-swedish",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "expectation for 2 pulls      = 1.0977443609022557\n",
      "From the generating function = 1.0977443609022557\n"
     ]
    }
   ],
   "source": [
    "# Checking a simple case\n",
    "# - expectation for 2 pulls longhand\n",
    "# - then extract from the matrix of coeffs to check\n",
    "\n",
    "grid = np.zeros((11,11))\n",
    "\n",
    "for i,j in itertools.product(range(11),range(11)):\n",
    "    grid[i,j] = (i+2)/77*(j+2-(i==j))/76\n",
    "    \n",
    "print(\"expectation for 2 pulls      =\",1+np.sum(grid * np.eye(11)))\n",
    "\n",
    "x = perms[0,-3] # ways of pulling 2 with a cap of 1 => x/y is expectation of 1\n",
    "y = perms[1,-3] # ways of pulling 2 with a cap of 2 => (y-x)/y is expectation of 2\n",
    "\n",
    "print(\"From the generating function =\",((2*(y-x)+x)/y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "signal-fault",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expectation is : 6.859787\n"
     ]
    }
   ],
   "source": [
    "# Generalizing the above for the whole matrix\n",
    "# lopping off the last column as it represents 0 draws\n",
    "# 1 is equally likely to be drawn at every point so divide by 78\n",
    "\n",
    "processed = np.zeros((12,78))\n",
    "processed[0,:]=perms[0,:]\n",
    "for i in range(11,0,-1):\n",
    "    processed[i,:] = perms[i,:] - perms [i-1,:]\n",
    "processed = (processed/np.sum(processed,axis=0))[:,:77]*np.array([np.arange(1,13)]).T\n",
    "\n",
    "print(\"Expectation is : {:.6f}\".format(np.sum(processed)/78))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "guilty-survey",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This month’s puzzle required a careful computation to remain tractable.  One approach that would work: by keeping track of the position of the partridge figurine and the maximum count of identical figurines drawn before it as the twelve different sets of figurines were inserted into the permutation reduced the space into a matrix of size 78 x 12.  The final expected value was approximately 6.859787\n"
     ]
    }
   ],
   "source": [
    "url='https://www.janestreet.com/puzzles/solutions/january-2021-solution/'\n",
    "res = requests.get(url)\n",
    "soup = BeautifulSoup(res.content, 'html.parser')\n",
    "x =[text for text in soup.body.stripped_strings]\n",
    "\n",
    "print(\" \".join(x[7:9]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "facial-airfare",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.0\n",
      "6.85959514040486\n",
      "6.859817070091465\n",
      "6.858856713714429\n",
      "6.859196785200804\n",
      "6.858591228281754\n",
      "6.858106190315635\n",
      "6.859182305831099\n",
      "6.859125392609326\n",
      "6.859076237880418\n",
      "\n",
      " Solution took 433.4376 seconds\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Check with a simulation to see if I'm close\n",
    "\n",
    "urn = np.array(\n",
    "      [1,\n",
    "       2,2,\n",
    "       3,3,3,\n",
    "       4,4,4,4,\n",
    "       5,5,5,5,5,\n",
    "       6,6,6,6,6,6,\n",
    "       7,7,7,7,7,7,7,\n",
    "       8,8,8,8,8,8,8,8,\n",
    "       9,9,9,9,9,9,9,9,9,\n",
    "       10,10,10,10,10,10,10,10,10,10,\n",
    "       11,11,11,11,11,11,11,11,11,11,11,\n",
    "       12,12,12,12,12,12,12,12,12,12,12,12],dtype=int)\n",
    "\n",
    "def loop():\n",
    "    temp = copy.copy(urn)\n",
    "    np.random.shuffle(temp)\n",
    "    picks = np.zeros(11)\n",
    "    for draw in temp:\n",
    "        if draw == 1:\n",
    "            return np.max(picks)\n",
    "        else:\n",
    "            picks[draw-2] +=1\n",
    "            \n",
    "start = time.perf_counter()\n",
    "sum = 0\n",
    "for i in range(10000000):\n",
    "    sum += loop()\n",
    "    if i % 1000000 ==0:\n",
    "        print(sum/(i+1))\n",
    "        \n",
    "stop =  time.perf_counter()\n",
    "print('\\n Solution took {:0.4f} seconds\\n'.format((stop-start)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "greek-cologne",
   "metadata": {},
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}