ooh maybe handling of empty and full set probabilities?

Author: R. Teal Witter

src/shapiq/approximator/sampling.py

@@ -237,7 +237,7 @@ def sample(self, budget: int):
         budget += budget % 2
 
         # Get sampling probabilities
-        self.get_scale_for_sampling(budget-2) # minus 2 for empty and full coalitions
+        self.get_scale_for_sampling(budget-2) # Exclude empty and full coalitions from budget
         sizes = np.arange(1, self.n_players)
         samples_per_size = self.symmetric_round_even(
             self.get_sampling_probs(sizes) * binom(self.n_players, sizes)
@@ -329,8 +329,11 @@ def coalitions_probability(self) -> np.ndarray:
         Returns:
             A copy of the sampled coalitions probabilities of shape ``(n_coalitions,)``
         """
-        return self.get_sampling_probs(self.coalitions_size)
-
+        probs = self.get_sampling_probs(self.coalitions_size)
+        # Replace the empty and full coalition probabilities with 1
+        probs[self.empty_coalition_index] = 1.0
+        probs[self.full_coalition_index] = 1.0
+        return probs
 
     @property
     def sampling_adjustment_weights(self) -> np.ndarray:
@@ -365,6 +368,19 @@ def empty_coalition_index(self) -> int | None:
         except IndexError:
             pass
         return None
+
+    @property
+    def full_coalition_index(self) -> int | None:
+        """
+        Returns:
+            The index of the full coalition or ``None`` if the full coalition was not sampled.
+        """
+        try:
+            if self.coalitions_per_size[-1] >= 1:
+                return int(np.where(self.coalitions_size == self.n_players)[0][0])
+        except IndexError:
+            pass
+        return None
 
     def set_random_state(self, random_state: int | None) -> None:
         '''