Day11_improved

day11_improved.py

import pandas as pd
import numpy as np

from numba import njit

def get_a(data):
    d = int(data)
    j = np.arange(1, 301)
    i = np.arange(1, 301)
    r = j + 10
    p = (r * i[:, None] + d) * r // 100 % 10 - 5
    return p

@njit
def f(c):
  n = c.shape[0]
  n = n - 1
  maxpower = np.array([0, 0, 0, 0])
  
  for i in range(n):
    for j in range(n):
      for k in range(n - max(i, j)):
        power = g(c, i, j, k)
        if power > maxpower[-1]:
          maxpower[:] = [j + 1, i + 1, k, power]
          
  return maxpower
        
@njit
def g(a, i, j, k):
  return a[i+k, j+k] - a[i, j+k] - a[i+k, j] + a[i, j]
        
def pad(a):
  n, m = a.shape
  c = np.zeros((n + 1, m + 1), a.dtype)
  c[1:, 1:] = a.cumsum(0).cumsum(1)
  
  return c

def part2(data):
  *a, p = f(pad(get_a(int(data))))
  return f"""Power: {p}
  AoC: {','.join(map(str, a))}"""

print(part2(18))
print(part2(42))
print(part2(5177))
%timeit part2(5177)