1

Can anyone suggest a quick and efficient solution for splitting an array of N elements into K groups with monotonicity maintaining? For example:

mas = [1, 2, 3, 4, 5]
K = 3

Expected output:

[1], [2], [3, 4, 5]
[1, 2], [3], [4, 5]
[1, 2, 3], [4], [5]
[1, 2], [3, 4], [5]
[1], [2, 3], [4, 5]
[1], [2, 3, 4], [5]

I can solve this problem by exhaustive search and a large number of checks, but what if someone has already solved a similar problem and can suggest a more efficient way or some idea to start?

BrokenBenchmark
  • 18,126
  • 7
  • 21
  • 33
Alex Lis
  • 53
  • 4

2 Answers2

2

Basically, you need to figure out combinations for choosing K-1 locations from total possible locations which is len(mas) - 1.

import itertools


# Data
mas = [1, 2, 3, 4, 5]
K = 3

ans = []
all_pos = list(range(1, len(mas)))
for pos in itertools.combinations(all_pos, K - 1):
    tmp = []
    prev_idx = 0
    for curr_idx in pos:
        tmp.append(mas[prev_idx:curr_idx])
        prev_idx = curr_idx
    tmp.append(mas[curr_idx:])
    ans.append(tmp)

ans
# [[[1], [2], [3, 4, 5]],
#  [[1], [2, 3], [4, 5]],
#  [[1], [2, 3, 4], [5]],
#  [[1, 2], [3], [4, 5]],
#  [[1, 2], [3, 4], [5]],
#  [[1, 2, 3], [4], [5]]]
d.b
  • 32,245
  • 6
  • 36
  • 77
2

If we can generate all possible valid lengths for each list, then we can generate these lengths and then read that out into our result using list slicing.

More specifically, we need to generate all possible permutations of positive integers such that they sum to the length of the list to partition.

A minor variant of this question happened to be asked on StackOverflow in 2011. (This variant merely required that all the numbers be non-negative, rather than strictly positive.) Using the code from this answer with minor modifications, we can generate the desired permutations and then read them off into our final result:

def sums(length, total_sum):
    if total_sum == 0:
        return
    if length == 1:
        yield (total_sum,)
    else:
        for value in range(1, total_sum + 1):
            for permutation in sums(length - 1, total_sum - value):
                yield (value,) + permutation
               
mas = [1, 2, 3, 4, 5]
K = 3

results = []
for lengths in sums(K, len(mas)):
    slices = []
    start = 0
    for slice_length in lengths:
        slices.append(mas[start:start+slice_length])
        start += slice_length
    results.append(slices)
print(results)

You can make this code more concise by generating the start index using accumulate() and chain() from itertools:

mas = [1, 2, 3, 4, 5]
K = 3

results = []
for lengths in sums(K, len(mas)):
    slices = []
    for start, slice_length in zip(chain([0], accumulate(lengths)), lengths):
        slices.append(mas[start:start+slice_length])
    results.append(slices)

These output:

[
 [[1], [2], [3, 4, 5]],
 [[1], [2, 3], [4, 5]],
 [[1], [2, 3, 4], [5]],
 [[1, 2], [3], [4, 5]],
 [[1, 2], [3, 4], [5]],
 [[1, 2, 3], [4], [5]]
]
BrokenBenchmark
  • 18,126
  • 7
  • 21
  • 33