Update app.py

app.py

@@ -23,17 +23,10 @@ Rk = np.indices((6,) * 6).reshape(6, -1).T
 Rk = Rk[Rk.sum(axis=1) <= 5]  # size 462
 R5 = Rk[Rk.sum(axis=1) == 5]  # size 252
 
-# Compute R (all combinations of up to 5 dice)
-R = np.indices((7,) * 5).reshape(5, -1).T
-R = np.array([np.bincount(r, minlength=7)[1:] for r in R])
-R = [R[R.sum(axis=1) == i] for i in range(6)]
-
-# Calculate probabilities (P) of each R5 outcome from each Rk roll
-P = np.zeros((len(Rk), len(R5)))
-R5_indices = {tuple(r): i for i, r in enumerate(R5)}
-for i, rk in enumerate(Rk):
-    for r in R[5 - rk.sum()]:
-        P[i, R5_indices[tuple(rk + r)]] += 1
+# Calculate probabilities (P) of each R5 outcome from each Rk roll using multinomial weights
+fact = np.array([1, 1, 2, 6, 24, 120])
+Δ = R5[None, :] - Rk[:, None]  # Δ[rk, r5, i]: number of die #i missing to go from Rk[rk] to R5[r5]
+P  = np.where((Δ >= 0).all(axis=2), fact[Δ.sum(axis=2)] // np.prod(fact[Δ.clip(min=0)], axis=2), 0)
 P /= P.sum(axis=1, keepdims=True)
 M = (P > 0).T  # Mask for valid transitions