enhance _iter_comparable loop

ChareefWong · Jul 3, 2024 · 97bcb63 · 97bcb63
1 parent 03566f6
commit 97bcb63
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/sksurv/metrics.py b/sksurv/metrics.py
@@ -96,12 +96,14 @@ def _iter_comparable(event_indicator, event_time, order):
         # check for tied event times
         event_at_same_time = event_indicator[order[i:end]]
         censored_at_same_time = ~event_at_same_time
+
+        mask = np.zeros(n_samples, dtype=bool)
+        mask[end:] = True
+        # an event is comparable to censored samples at same time point
+        mask[i:end] = censored_at_same_time
+
         for j in range(i, end):
             if event_indicator[order[j]]:
-                mask = np.zeros(n_samples, dtype=bool)
-                mask[end:] = True
-                # an event is comparable to censored samples at same time point
-                mask[i:end] = censored_at_same_time
                 tied_time += censored_at_same_time.sum()
                 yield (j, mask, tied_time)
         i = end