Solution to problem 11 in Python

src/Year_2018/P11.py

@@ -90,23 +90,23 @@ def get_value(x: int, y: int) -> int:
     return power_level - 5
 
 
+partial_sums_p1: dict[tuple[int, int], int] = defaultdict(int)
+
+
 def calculate_ps(x, y):
     return (
         get_value(x + 1, y + 1)
-        + partial_sums[x, y - 1]
+        + partial_sums_p1[x, y - 1]
-        + partial_sums[x - 1, y]
+        + partial_sums_p1[x - 1, y]
-        - partial_sums[x - 1, y - 1]
+        - partial_sums_p1[x - 1, y - 1]
     )
 
 
-partial_sums = defaultdict(int)
-
-
 def get_partial_sums(arr_size: int) -> dict[tuple[int, int], int]:
     for y in range(arr_size):
         for x in range(arr_size):
-            partial_sums[(x, y)] = calculate_ps(x, y)
+            partial_sums_p1[(x, y)] = calculate_ps(x, y)
-    return partial_sums
+    return partial_sums_p1
 
 
 def part_1() -> None:
@@ -130,5 +130,49 @@ def part_1() -> None:
     print(f"Largest total power is located at {x3},{y3}")
 
 
+# --- Part Two ---
+
+# You discover a dial on the side of the device; it seems to let you select a
+# square of any size, not just 3x3. Sizes from 1x1 to 300x300 are supported.
+
+# Realizing this, you now must find the square of any size with the largest
+# total power. Identify this square by including its size as a third parameter
+# after the top-left coordinate: a 9x9 square with a top-left corner of 3,5 is
+# identified as 3,5,9.
+
+# For example:
+
+#     For grid serial number 18, the largest total square (with a total power
+# of 113) is 16x16 and has a top-left corner of 90,269, so its identifier is
+# 90,269,16.
+#     For grid serial number 42, the largest total square (with a total power
+# of 119) is 12x12 and has a top-left corner of 232,251, so its identifier is
+# 232,251,12.
+
+# What is the X,Y,size identifier of the square with the largest total power?
+
+
+def part_2() -> None:
+    arr_size = 300
+    grid_sums = {}
+
+    partial_sums = get_partial_sums(arr_size)
+
+    for size in range(1, arr_size):
+        for j in range(size - 1, arr_size):
+            for i in range(size - 1, arr_size):
+                gp = (
+                    partial_sums[(i, j)]
+                    + partial_sums[(i - size, j - size)]
+                    - partial_sums[(i - size, j)]
+                    - partial_sums[(i, j - size)]
+                )
+                grid_sums[gp] = (i - size + 2, j - size + 2, size)
+
+    x_max, y_max, size_max = grid_sums[max(grid_sums)]
+    print(f"Largest total power identifier is {x_max},{y_max},{size_max}")
+
+
 if __name__ == "__main__":
     part_1()
+    part_2()