5
\$\begingroup\$

This code should partition a list into two subsets of the given list that add up to the same number: for example [15, 5, 20, 10, 35, 25, 10] -> [15, 5, 10, 15, 10], [20, 35] would be a valid solution. I'm looking for any tips to improve code style (be more Pythonic) or performance (mainly improvements to the algorithm used, but Python-specific things are also welcome). Thanks!

from collections import Counter


def partition_into_equal_parts(l):
    '''Partitions s into two subsets of l that have the same sum.

    >>> problem = [15, 5, 20, 10, 35, 25, 10]
    >>> first, second = partition_into_equal_parts(problem)
    >>> valid_solution(first, second, problem)
    True
    '''
    total = sum(l)
    # If sum is odd, there is no way that total = sum(first) + sum(second) = 2 * sum(first)
    if total % 2:
        return
    first = subset_sum(total // 2, l)
    if first is None:
        return
    second = []
    # Fill second with items from counter
    second_counter = Counter(l) - Counter(first)
    for number, amount in second_counter.items():
        second.extend([number] * amount)
    return first, second


def valid_solution(first, second, problem):
    return sum(first) == sum(second) and Counter(first) + Counter(second) == Counter(problem)


def subset_sum(k, lst):
    '''Returns a subset of lst that has a sum of k.

    >>> sum(subset_sum(24, [12, 1, 61, 5, 9, 2]))
    24
    >>> subset_sum(53, [12, 13, 14])
    '''
    return recursive_calculate(k, sorted(lst, reverse=True), 0)


def recursive_calculate(k, lst, start):
    for idx in range(start, len(lst)):
        if lst[idx] == k:
            return [lst[idx]]
        elif lst[idx] < k:
            rest = recursive_calculate(k - lst[idx], lst, idx + 1)
            if rest is not None:
                rest.append(lst[idx])
                return rest
\$\endgroup\$
  • 2
    \$\begingroup\$ Good job on checking the early out if sum of the input is odd! In terms of algorithm, finding a subset of a list which totals half the total of the whole list is a special case of the knapsack problem. It's NP-complete in principle, but there are effective algorithms for some cases. In particular, look up the nice dynamic programming algorithm that works so long as the numbers are all smallish. \$\endgroup\$ – Josiah Apr 25 '18 at 19:02
  • \$\begingroup\$ @Josiah With a 2D table subset, the value of subset[i][j] will be true if there is a subset of set[0..j-1] with sum equal to i. This would indeed be pretty bad with larger numbers. It would be equivalent to replace it with a 2D table with subset[j][i] (same table, but rotated). Would it then be efficient to use subset[j]{i}, where {i} is a set, so you would have a 1D array of sets? \$\endgroup\$ – redfast00 Apr 25 '18 at 19:12
  • \$\begingroup\$ That I don't know. You'd have to profile it with your data. \$\endgroup\$ – Josiah Apr 25 '18 at 22:55
1
\$\begingroup\$

Some "more pythonic" suggestions which also improve performance:

  1. use Counter.elements()

in partition_into_equal_parts this code:

second = []
# Fill second with items from counter
second_counter = Counter(l) - Counter(first)
for number, amount in second_counter.items():
    second.extend([number] * amount)

can be merely replaced by

second = list((Counter(l) - Counter(first)).elements()

As Counter.elements returns a chained iterable of the expanded elements according to their count. So no need to do that with a loop and an extra list.

  1. Cache the indexed element

In recursive_calculate, you have to use the index in

for idx in range(start, len(lst)):

but why not storing lst[idx] in a local variable? It's repeated a lot of times, and accessing a list element is costly. Just cache this access since you're not writing in lst (well you're writing to the end of it, recursively, but lst[idx] isn't udpated because the index is low enough).

def recursive_calculate(k, lst, start):
    for idx in range(start, len(lst)):
        element = lst[idx]
        if element == k:
            return [element]
        elif element < k:
            rest = recursive_calculate(k - element, lst, idx + 1)
            if rest is not None:
                rest.append(element)
                return rest
    return None

Small note: python functions implicitly return None, but "Explicit is better than implicit" (The Zen Of Python) so it's better to actually return it.

|improve this answer|||||
\$\endgroup\$
  • \$\begingroup\$ Using Counter.elements was exactly what I was looking for, thanks! I also agree with the second point and the small note. I will leave this open for some time, but if I don't accept an answer after some time, ping me. \$\endgroup\$ – redfast00 Apr 26 '18 at 20:06

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.