From ed859d0f0845e921b991da0d8a3e11f3409ab24c Mon Sep 17 00:00:00 2001
From: tosemml <97384583+tosemml@users.noreply.github.com>
Date: Tue, 8 Aug 2023 12:14:34 -0700
Subject: [PATCH] Minor code  improvements

---
 query_strategies/util.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/query_strategies/util.py b/query_strategies/util.py
index 819cee4..bc55728 100644
--- a/query_strategies/util.py
+++ b/query_strategies/util.py
@@ -139,10 +139,7 @@ def interleave_offsets(batch, nu):
     groups = [batch // (nu + 1)] * (nu + 1)
     for x in range(batch - sum(groups)):
         groups[-x - 1] += 1
-    offsets = [0]
-    for g in groups:
-        offsets.append(offsets[-1] + g)
-    assert offsets[-1] == batch
+    offsets = [0] + np.cumsum(groups).tolist()
     return offsets
 
 def interleave(xy, batch):