Advertisement
kosievdmerwe

Untitled

Jul 29th, 2023
834
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.97 KB | None | 0 0
  1. from copy import deepcopy
  2.  
  3. MOVES = [
  4.     (4, 0),
  5.     (3, 1),
  6.     (2, 2),
  7.     (1, 3),
  8. ]
  9.  
  10.  
  11. class Matrix:
  12.     def __init__(self, rows: int, cols: int):
  13.         self.rows = rows
  14.         self.cols = cols
  15.         self.grid = [
  16.             [0] * cols
  17.             for _ in range(rows)
  18.         ]
  19.  
  20.     def __add__(self, other: "Matrix") -> "Matrix":
  21.         assert self.rows == other.rows
  22.         assert self.cols == other.cols
  23.  
  24.         ans = Matrix(self.rows, self.cols)
  25.         for r in range(self.rows):
  26.             for c in range(self.cols):
  27.                 ans.grid[r][c] = self.grid[r][c] + other.grid[r][c]
  28.         return ans
  29.  
  30.     def __mul__(self, other: "Matrix") -> "Matrix":
  31.         assert self.cols == other.rows
  32.  
  33.         rows = self.rows
  34.         cols = other.cols
  35.         ans = Matrix(rows, cols)
  36.         for r in range(rows):
  37.             for c in range(cols):
  38.                 ans.grid[r][c] = sum(
  39.                     self.grid[r][i] * other.grid[i][c]
  40.                     for i in range(self.cols)
  41.                 )
  42.         return ans
  43.  
  44.     def __pow__(self, power: int) -> "Matrix":
  45.         assert type(power) == int
  46.         assert power >= 1
  47.         assert self.rows == self.cols
  48.  
  49.         rows = self.rows
  50.  
  51.         ans = Matrix(rows, rows)
  52.         for i in range(rows):
  53.             ans.grid[i][i] = 1
  54.  
  55.         cur = deepcopy(self)
  56.  
  57.         i = 1
  58.         while i <= power:
  59.             if power & i:
  60.                 ans = cur * ans
  61.             cur = cur * cur
  62.             i <<= 1
  63.  
  64.         return ans
  65.  
  66.     def print(self) -> None:
  67.         for row in self.grid:
  68.             print(row)
  69.         print()
  70.  
  71.  
  72. class Solution:
  73.     def soupServings(self, n: int) -> float:
  74.         n = (n + 24) // 25  # ceil(n / 25)
  75.  
  76.         @cache
  77.         def prob(a: int, b: int) -> float:
  78.             if a <= 0 and b <= 0:
  79.                 return 0.5
  80.            
  81.             if a <= 0:
  82.                 return 1.0
  83.  
  84.             if b <= 0:
  85.                 return 0.0
  86.            
  87.             return sum(
  88.                 0.25 * prob(a - ad, b - bd)
  89.                 for ad, bd in MOVES
  90.             )
  91.  
  92.         deltas = [(ad, bd) for ad in range(5) for bd in range(5)]
  93.         D = len(deltas)
  94.         idxes = {d: i for i, d in enumerate(deltas)}
  95.  
  96.         A = Matrix(D, D)
  97.         B = Matrix(D, 1)
  98.  
  99.         for (ad, bd), idx in idxes.items():
  100.             B.grid[idx][0] = prob(4 - ad, 4 - bd)
  101.  
  102.         B.print()
  103.  
  104.         for ad, bd in deltas:
  105.             k = (5 + ad, 5 + bd)
  106.             idx = idxes[(k[0] - 1, k[1] - 1)]
  107.  
  108.             if k in idxes:
  109.                 A.grid[idx][idxes[k]] = 1
  110.                 continue
  111.  
  112.             for mad, mbd in MOVES:
  113.                 mk = (
  114.                     k[0] - mad,
  115.                     k[1] - mbd,
  116.                 )
  117.                 A.grid[idx][idxes[mk]] = 0.25
  118.  
  119.         A.print()
  120.        
  121.         if n <= 5:
  122.             return prob(n, n)
  123.  
  124.         C = A ** (n - 4) * B
  125.         return C.grid[idxes[(0, 0)]][0]
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement