summaryrefslogtreecommitdiff
path: root/utilities.py
diff options
context:
space:
mode:
Diffstat (limited to 'utilities.py')
-rw-r--r--utilities.py23
1 files changed, 23 insertions, 0 deletions
diff --git a/utilities.py b/utilities.py
new file mode 100644
index 0000000..1ad24b8
--- /dev/null
+++ b/utilities.py
@@ -0,0 +1,23 @@
+from typing import Tuple
+from random import shuffle
+
+
+def _shuffle_two_lists(a: list, b: list) -> Tuple[list, list]:
+ indices = list(range(len(a)))
+ shuffle(indices)
+ a = [a[index] for index in indices]
+ b = [b[index] for index in indices]
+ return a, b
+
+
+def _pick_from_list_after_sorting_other(a: list, b: list[int]) -> list:
+ indices = [i[0] for i in sorted(enumerate(b), key=lambda x: x[1])]
+ return [a[index] for index in indices]
+
+
+def sample_list(indices: list, cost: list[int], count: int = None) -> list:
+ if count is None:
+ count = len(indices)
+ indices, cost = _shuffle_two_lists(indices, cost)
+ indices = _pick_from_list_after_sorting_other(indices, cost)[:count]
+ return indices