diff --git a/backend/examples/assignment/strategies.py b/backend/examples/assignment/strategies.py index 22c7f1e0..f1eb85d5 100644 --- a/backend/examples/assignment/strategies.py +++ b/backend/examples/assignment/strategies.py @@ -45,8 +45,8 @@ class WeightedSequentialStrategy(BaseStrategy): def assign(self) -> List[Assignment]: assignments = [] - ratio = list(np.round(np.cumsum(self.weights) / 100 * self.dataset_size).astype(int)) - ratio = [0] + ratio[:-1] + [self.dataset_size] + cumsum = np.cumsum([0] + self.weights) + ratio = np.round(cumsum / 100 * self.dataset_size).astype(int) for user, (start, end) in enumerate(zip(ratio, ratio[1:])): # Todo: use itertools.pairwise assignments.extend([Assignment(user=user, example=example) for example in range(start, end)]) return assignments