From 464a931df3eea396a179c396f6826163ab6b6e9e Mon Sep 17 00:00:00 2001
From: Gurge44 <gurge44.minecraft@gmail.com>
Date: Sat, 3 Jan 2026 11:33:02 +0100
Subject: [PATCH] improve weighted role pick

---
 Modules/CustomRoleSelector.cs    | 8 ++++++--
 Roles/Crewmate/Retributionist.cs | 1 +
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/Modules/CustomRoleSelector.cs b/Modules/CustomRoleSelector.cs
index 47f7caf23..51700df6c 100644
--- a/Modules/CustomRoleSelector.cs
+++ b/Modules/CustomRoleSelector.cs
@@ -438,7 +438,11 @@ internal static class CustomRoleSelector
         
         static RoleAssignInfo PickWeighted(List<RoleAssignInfo> pool, IRandom rng)
         {
-            int totalWeight = pool.Sum(t => t.SpawnChance);
+            int totalWeight = 0;
+
+            foreach (var info in pool)
+                totalWeight += info.SpawnChance * (info.MaxCount - info.AssignedCount);
+
             if (totalWeight <= 0) return null;
 
             int roll = rng.Next(totalWeight);
@@ -446,7 +450,7 @@ internal static class CustomRoleSelector
 
             foreach (var info in pool)
             {
-                cumulative += info.SpawnChance;
+                cumulative += info.SpawnChance * (info.MaxCount - info.AssignedCount);
                 if (roll < cumulative) return info;
             }
 
diff --git a/Roles/Crewmate/Retributionist.cs b/Roles/Crewmate/Retributionist.cs
index 2ea86370d..35f1d6ff5 100644
--- a/Roles/Crewmate/Retributionist.cs
+++ b/Roles/Crewmate/Retributionist.cs
@@ -2,6 +2,7 @@ using System;
 using System.Linq;
 using AmongUs.GameOptions;
 using EHR.Modules;
+using EHR.Neutral;
 using Hazel;
 using UnityEngine;