In [14]:

    

import numpy as np
from itertools import product
from tqdm import tqdm


    

In [15]:

    

serial = 1308


    

In [16]:

    

def power(x, y, serial):

    rackid = (x + 1) + 10
    power = rackid * (y + 1)
    power += serial
    power *= rackid
    return (abs(power) // 100) % 10 - 5


    

In [17]:

    

def best_square(serial, dim=300):
    
    grid = np.empty((300, 300))
    for i, j in product(range(300), repeat=2):
        grid[i, j] = power(i, j, serial)
    m = 0
    for (x, y), indices in square_generator():
        s = np.sum([grid[a] for a in indices])
        if s > m:
            a, m = (x, y), s
    return a, m
    
def square_generator():
    
    for i, j in product(range(300), repeat=2):
        if (300 - i - 1 >= 2) and (300 - j - 1 >= 2):
            yield (i, j), [(i + s, j + t) for s, t in product([0, 1, 2], repeat=2)]


    

In [25]:

    

def test():
    assert(power(121, 78, 57) == -5)
    assert(power(216, 195, 39) == 0)
    assert(power(100, 152, 71) == 4)

test()


    

In [26]:

    

(a, b), charge = best_square(serial)
print(str(a+1) + ',' + str(b+1))


    

    




21,41

In [27]:

    

class PowerGrid:
    
    def __init__(self, serial):
        self.grid = np.zeros((300, 300, 300))
        for i, j in product(range(300), repeat=2):
            self.grid[i, j, 0] = power(i, j, serial)
    
    def lshape(self, i, j, size):
        lateral = sum(self.grid[i: i + size, j + size, 0]) 
        bottom  = sum(self.grid[i + size, j: j + size, 0])
        corner  = self.grid[i + size, j + size, 0]
        return lateral + bottom + corner
    
    def dichotosum(self, i, j, q):
        return self.grid[i,j,q-1] + self.grid[i+q,j,q-1] + self.grid[i,j+q,q-1] + self.grid[i+q,j+q,q-1]
    
    def topsquare(self):
        msize, mx, index = 1, np.max(self.grid[:, :, 0]), np.argmax(self.grid[:, :, 0])
        size  = 2
        while size <= 300:
            for i, j in product(range(300), repeat=2):
                if (300 - i >= size) and (300 - j >= size):
                    if size % 2 == 0:
                        q = size // 2
                        a = 0
                    else:
                        q = (size - 1) // 2
                        a = self.lshape(i, j, 2*q)
                    a += self.dichotosum(i, j, q)
                    self.grid[i, j, size-1] = a
                    if a > mx:
                        msize   = size
                        index   = (i, j)
                        mx      = a
                        aoc_out = str(i+1) + ',' + str(j+1) + ',' + str(msize)
            size += 1
        return aoc_out


    

In [28]:

    

x = PowerGrid(18)
x.topsquare()


    

    
Out[28]:





'90,269,16'

In [30]:

    

y = PowerGrid(42)
y.topsquare()


    

    
Out[30]:





'232,251,12'

In [31]:

    

serial = 1308
z = PowerGrid(serial)
z.topsquare()


    

    
Out[31]:





'227,199,19'