import numpy as np
input = 4455
fuel_cells = np.zeros((301,301), dtype=np.int64)

for y in range(1, 301):
  fuel_cells[1:,y] = y
for x in range(1, 301):
  fuel_cells[x,1:] *= x+10
fuel_cells[1:,1:] += input
for x in range(1, 301):
  fuel_cells[x,1:] *= x+10
fuel_cells[1:,1:] //= 100
fuel_cells[1:,1:] %= 10
fuel_cells[1:,1:] -= 5
fuel_cs = fuel_cells.cumsum(axis=0).cumsum(axis=1)


def sum_square(x, y, s):
  A = fuel_cs[x-1, y-1]
  B = fuel_cs[x+s-1, y-1]
  C = fuel_cs[x-1, y+s-1]
  D = fuel_cs[x+s-1, y+s-1]
  return D - B - C + A


max_power_3 = 0
max_idx_3 = (0, 0, 0)
for x in range(1, 299):
  for y in range(1, 299):
    val = sum_square(x, y, 3)
    if val > max_power_3:
      max_power_3 = val
      max_idx_3 = (x, y, 3)

print(max_idx_3)  # Part 1


max_power_s = 0
max_idx_s = (0,0,0)
for x in range(1, 298):
  for y in range(1, 298):
    A = fuel_cs[x-1, y-1]
    Ds = fuel_cs[x+3:, y+3:].diagonal()
    Bs = fuel_cs[x+3:, y-1][:len(Ds)]
    Cs = fuel_cs[x-1, y+3:][:len(Ds)]
    sums = Ds - Cs - Bs + A
    idx = sums.argmax()
    val = sums[idx]
    if val > max_power_s:
      max_power_s = val
      max_idx_s = (x, y, 4+idx)

print(max_idx_s)  # Part 2