04.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from utils.get_input import get_input\n",
    "import numpy as np\n",
    "from numba import njit\n",
    "\n",
    "# from numba import jit\n",
    "\n",
    "# inp = Path(\"./test_input\").read_text()\n",
    "inp = get_input(4)\n",
    "inp;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.68 ms ± 7.18 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "lines = inp.strip().split(\"\\n\")\n",
    "inp_np = np.array([list(line) for line in lines])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2397"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@njit\n",
    "def find_xmas(x: np.array) -> int:\n",
    "    sum = 0\n",
    "    for i in range(len(x) - 3):\n",
    "        if x[i] == \"X\" and x[i + 1] == \"M\" and x[i + 2] == \"A\" and x[i + 3] == \"S\":\n",
    "            sum += 1\n",
    "    return sum\n",
    "\n",
    "\n",
    "@njit\n",
    "def part1():\n",
    "    sum = 0\n",
    "    # loop over all rows\n",
    "    for row in inp_np:\n",
    "        sum += find_xmas(row)\n",
    "\n",
    "        # Reverse\n",
    "        row_rev = row[::-1]\n",
    "        sum += find_xmas(row_rev)\n",
    "\n",
    "    # loop over all columns\n",
    "    for col in inp_np.T:\n",
    "        sum += find_xmas(col)\n",
    "\n",
    "        # Reverse\n",
    "        col_rev = col[::-1]\n",
    "        sum += find_xmas(col_rev)\n",
    "\n",
    "    # Loop over all diagonals\n",
    "    for i in range(-len(inp_np) + 1, len(inp_np)):\n",
    "\n",
    "        ## Diag 1\n",
    "        # x\n",
    "        #  x\n",
    "        #   x\n",
    "        diag = np.diag(inp_np, i)\n",
    "        sum += find_xmas(diag)\n",
    "\n",
    "        # Reverse\n",
    "        diag_rev = diag[::-1]\n",
    "        sum += find_xmas(diag_rev)\n",
    "\n",
    "        ## Diag 2\n",
    "        #   x\n",
    "        #  x\n",
    "        # x\n",
    "        diag = np.diag(np.fliplr(inp_np), i)\n",
    "        sum += find_xmas(diag)\n",
    "\n",
    "        # Reverse\n",
    "        diag_rev = diag[::-1]\n",
    "        sum += find_xmas(diag_rev)\n",
    "    return sum\n",
    "\n",
    "\n",
    "\n",
    "part1()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.52 ms ± 60.7 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "part1()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1824"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@njit\n",
    "def find_mas(x: np.array) -> int:\n",
    "    if x[0] == \"M\" and x[1] == \"A\" and x[2] == \"S\":\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "\n",
    "@njit\n",
    "def part2():\n",
    "    sum = 0\n",
    "    # Loop over all diagonals\n",
    "    for x in range(inp_np.shape[0] - 2):\n",
    "        for y in range(inp_np.shape[1] - 2):\n",
    "\n",
    "            sub = inp_np[x : x + 3, y : y + 3]\n",
    "            diag = np.diag(sub, 0)\n",
    "\n",
    "            if find_mas(diag) or find_mas(diag[::-1]):\n",
    "                diag_flip = np.diag(np.fliplr(sub), 0)\n",
    "                if find_mas(diag_flip) or find_mas(diag_flip[::-1]):\n",
    "                    sum += 1\n",
    "    return sum\n",
    "\n",
    "\n",
    "\n",
    "part2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6 ms ± 102 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "part2()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}