day_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "import polars.selectors as cs\n",
    "\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[WindowsPath('../data/input_1.txt'), WindowsPath('../data/input_2.txt')]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATA_PATH = Path('..', 'data')\n",
    "list(DATA_PATH.glob('*'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_MAP = {\n",
    "    'red': 12,\n",
    "    'green': 13,\n",
    "    'blue': 14,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = (\n",
    "    pl\n",
    "    .read_csv(DATA_PATH / 'input_2.txt', has_header=False, separator='\\n')\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id_game</th><th>id_set</th><th>blue</th><th>red</th><th>green</th></tr><tr><td>i32</td><td>i32</td><td>i64</td><td>i64</td><td>i64</td></tr></thead><tbody><tr><td>1</td><td>1</td><td>12</td><td>15</td><td>2</td></tr><tr><td>1</td><td>2</td><td>5</td><td>17</td><td>8</td></tr><tr><td>1</td><td>3</td><td>17</td><td>8</td><td>0</td></tr><tr><td>1</td><td>4</td><td>1</td><td>4</td><td>9</td></tr><tr><td>2</td><td>1</td><td>6</td><td>6</td><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 5)\n",
       "┌─────────┬────────┬──────┬─────┬───────┐\n",
       "│ id_game ┆ id_set ┆ blue ┆ red ┆ green │\n",
       "│ ---     ┆ ---    ┆ ---  ┆ --- ┆ ---   │\n",
       "│ i32     ┆ i32    ┆ i64  ┆ i64 ┆ i64   │\n",
       "╞═════════╪════════╪══════╪═════╪═══════╡\n",
       "│ 1       ┆ 1      ┆ 12   ┆ 15  ┆ 2     │\n",
       "│ 1       ┆ 2      ┆ 5    ┆ 17  ┆ 8     │\n",
       "│ 1       ┆ 3      ┆ 17   ┆ 8   ┆ 0     │\n",
       "│ 1       ┆ 4      ┆ 1    ┆ 4   ┆ 9     │\n",
       "│ 2       ┆ 1      ┆ 6    ┆ 6   ┆ 2     │\n",
       "└─────────┴────────┴──────┴─────┴───────┘"
      ]
     },
     "execution_count": 215,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_game_set = (\n",
    "    df\n",
    "    .with_columns(\n",
    "        pl.col('column_1').str.split(': '),\n",
    "    )\n",
    "    .select(\n",
    "        pl.col('column_1').list[0].str.replace('Game ', '').cast(pl.Int32).alias('id_game'),\n",
    "        pl.col('column_1').list[1].alias('sets').str.split('; ')\n",
    "    )\n",
    "    .explode('sets')\n",
    "    .with_columns(\n",
    "        pl.lit(1).alias('id_set')\n",
    "    )\n",
    "    .with_columns(\n",
    "        pl.col('id_set').cum_sum().over('id_game')\n",
    "    )\n",
    "    .with_columns(\n",
    "        pl.col('sets').str.split(', ')\n",
    "    )\n",
    "    .explode('sets')\n",
    "    .with_columns(\n",
    "        pl.col('sets').str.extract_groups('(\\d+) ([a-zA-Z]+)'),\n",
    "    )\n",
    "    .unnest('sets')\n",
    "    .rename({'1': 'value', '2': 'color'})\n",
    "    .with_columns(\n",
    "        pl.col('value').cast(pl.Int64),\n",
    "    )\n",
    "    .pivot(\n",
    "        index=['id_game', 'id_set'],\n",
    "        columns='color',\n",
    "        values='value',\n",
    "    )\n",
    "    .fill_null(0)\n",
    ")\n",
    "df_game_set.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id_game</th><th>is_valid</th></tr><tr><td>i32</td><td>u32</td></tr></thead><tbody><tr><td>2716</td><td>49</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 2)\n",
       "┌─────────┬──────────┐\n",
       "│ id_game ┆ is_valid │\n",
       "│ ---     ┆ ---      │\n",
       "│ i32     ┆ u32      │\n",
       "╞═════════╪══════════╡\n",
       "│ 2716    ┆ 49       │\n",
       "└─────────┴──────────┘"
      ]
     },
     "execution_count": 216,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    df_game_set\n",
    "    .pipe(\n",
    "        lambda x: (\n",
    "            x\n",
    "            .with_columns([\n",
    "                pl.col(color).le(MAX_MAP[color])\n",
    "                for color in x.select(~cs.contains('id')).columns\n",
    "            ])\n",
    "            .with_columns(\n",
    "                pl\n",
    "                .fold(\n",
    "                    acc=True,\n",
    "                    function=lambda s1, s2: s1 & s2,\n",
    "                    exprs=~cs.contains('id')\n",
    "                )\n",
    "                .alias('is_valid')\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "    .group_by('id_game', maintain_order=True)\n",
    "    .agg(pl.col('is_valid').all())\n",
    "    .filter(pl.col('is_valid'))\n",
    "    .sum()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 240,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "72227"
      ]
     },
     "execution_count": 240,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    df_game_set\n",
    "    .group_by('id_game', maintain_order=True)\n",
    "    .agg(\n",
    "        pl.all().max()\n",
    "    )\n",
    "    .select(\n",
    "        ~cs.contains('id')\n",
    "    )\n",
    "    .fold(\n",
    "        lambda s1, s2: s1 * s2\n",
    "    )\n",
    "    .sum()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pruebas2",
   "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}