{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "beginning-smoke",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class Board():    \n",
    "    def __init__(self, size):\n",
    "        self.numbers = np.zeros((size, size), np.int32)\n",
    "        self.marks = np.zeros((size, size), bool)\n",
    "        \n",
    "    def is_winner(self):\n",
    "        size = self.marks.shape[0]\n",
    "        for i in range(size):\n",
    "            if np.sum(self.marks[i]) == size or \\\n",
    "                np.sum(self.marks[:,i]) == size:\n",
    "                \n",
    "                return True\n",
    "    \n",
    "    def mark(self, n):\n",
    "        self.marks |= self.numbers == n\n",
    "    \n",
    "    def sum_unmarked(self):\n",
    "        return np.sum(self.numbers * ~self.marks)\n",
    "\n",
    "def parse_file(f):\n",
    "    numbers_drawn = map(int, next(f).split(','))\n",
    "    board = None\n",
    "    boards = []\n",
    "    \n",
    "    i = 0\n",
    "    \n",
    "    for l in f:\n",
    "        l = l.rstrip()\n",
    "        if l == \"\":\n",
    "            if board:\n",
    "                boards.append(board)\n",
    "            board = Board(5)\n",
    "            i = 0\n",
    "            continue\n",
    "            \n",
    "        board.numbers[i] = list(map(int, l.split()))\n",
    "        i += 1\n",
    "        \n",
    "    return numbers_drawn, boards\n",
    "\n",
    "def find_winner(numbers_drawn, boards):\n",
    "    first_winner_score = None\n",
    "    last_winner_score = None\n",
    "    numbers_drawn = list(numbers_drawn)\n",
    "    for ni, n in enumerate(numbers_drawn):\n",
    "        for b in boards:\n",
    "            b.mark(n)\n",
    "            \n",
    "        for bi, b in enumerate(boards):\n",
    "            if b.is_winner():\n",
    "                boards[bi] = None\n",
    "                \n",
    "                if first_winner_score is None:\n",
    "                    first_winner_score = b.sum_unmarked() * n\n",
    "                    \n",
    "                if len(boards) == 1:\n",
    "                    last_winner_score  = b.sum_unmarked() * n\n",
    "\n",
    "        boards = [b for b in boards if b]\n",
    "                    \n",
    "    return first_winner_score, last_winner_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "curious-honduras",
   "metadata": {},
   "outputs": [],
   "source": [
    "import import_ipynb\n",
    "import helper\n",
    "f = helper.open_file('2021_4.txt')\n",
    "numbers_drawn, boards = parse_file(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "respected-testament",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(64084, 12833)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "find_winner(numbers_drawn, boards)"
   ]
  }
 ],
 "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.9.2"
  },
  "title": "Giant Squid"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}