def assign_levels(super_id, indices, levels):
if not isinstance(indices, list):
raise TypeError('indices is invalid')
ret = []
if not isinstance(indices, list) or len(indices) == 0:
return ret
if isinstance(levels, list) and len(levels) > 0:
current = levels[0]
levels = levels[1:]
if len(levels) == 0:
for i in indices:
item = dict(
level=current['level'],
index=i,
super=super_id
)
ret.append(item)
else:
# determines the number of branches
if 'count' in current:
_count = current['count']
else:
_min = current['min'] if 'min' in current else 0
_max = current['max'] if 'max' in current else len(indices)
_count = random.randint(_min, _max)
# no any branch of this level
if _count == 0:
return assign_levels(super_id, indices, levels)
# split the children
group_len = _count if _count < len(indices) else len(indices)
child_group = array_random_split(indices, group_len)
for child in child_group:
first_index = child[0]
subchild = child[1:]
ret.append(dict(
level=current['level'],
index=first_index,
super=super_id
))
ret = ret + assign_levels(first_index, subchild, levels = levels)
return ret
def test_array_random_split_ok(self):
_arr = range(20)
_chunks = array_random_split(_arr, 5)
self.assertEqual(len(_chunks), 5)
def test_array_random_split_with_number_of_subarrays_less_than_array_size(self):
_arr = [1, 2, 3, 4]
with self.assertRaises(ValueError):
array_random_split(_arr, 5, empty_accepted=False)
def test_array_random_split_with_number_of_subarrays_equals_to_array_size(self):
_arr = [1, 2, 3, 4, 5]
for _ in range(3):
_chunks = array_random_split(_arr, 5, empty_accepted=False)
self.assertEqual(len(_chunks), 5)
self.assertEqual(_chunks, [[1], [2], [3], [4], [5]])
def test_array_random_split_empty_array_to_5_subarrays(self): _arr = [] _chunks = array_random_split(_arr, 3, empty_accepted=True) self.assertEqual(len(_chunks), 3) self.assertEqual(list(itertools.chain.from_iterable(_chunks)), _arr)
def test_array_random_split_with_number_of_subarrays_greater_than_array_size(self): _arr = [1, 2, 3] _chunks = array_random_split(_arr, 5, empty_accepted=True) self.assertEqual(len(_chunks), 5) self.assertEqual(list(itertools.chain.from_iterable(_chunks)), _arr)