1 files changed, 2 insertions, 11 deletions
diff --git a/6np.py b/6np.py
index 92273b1..8d90f1f 100644
--- a/6np.py
+++ b/6np.py
@@ -4,17 +4,8 @@ import numpy as np
 
 fish = Counter(map(int, open_day(6).read().split(',')))
 fish = np.array([fish[i] for i in range(9)])
-mat = np.array([
-    [0, 0, 0, 0, 0, 0, 1, 0, 1],
-    [1, 0, 0, 0, 0, 0, 0, 0, 0],
-    [0, 1, 0, 0, 0, 0, 0, 0, 0],
-    [0, 0, 1, 0, 0, 0, 0, 0, 0],
-    [0, 0, 0, 1, 0, 0, 0, 0, 0],
-    [0, 0, 0, 0, 1, 0, 0, 0, 0],
-    [0, 0, 0, 0, 0, 1, 0, 0, 0],
-    [0, 0, 0, 0, 0, 0, 1, 0, 0],
-    [0, 0, 0, 0, 0, 0, 0, 1, 0]
-])
+mat = np.eye(9, 9, -1, dtype=int)
+mat[0, [6, 8]] = 1
 
 print(sum(np.matmul(fish, np.linalg.matrix_power(mat, 80))))
 print(sum(np.matmul(fish, np.linalg.matrix_power(mat, 256))))