Search code examples
pythonarrayslist-comprehensioncombinationspermutation

In how many ways can I distribute "n" candies to 3 children such that each gets at least 1 candy and at least 1 child should get more than others


How many different ways can I distribute "n" candies to m = 3 children such that each one gets at least 1 candy and at least 1 child should get more candies than others?

For eg: No. of candies = 6, Total ways = 9
[1,4,1] , [1,1,4] , [4,1,1] , [1,2,3] , [1,3,2] and so on.... total ways = 9.

My code in Python:

import itertools

n = int(input())
arr = []

for i in range(1,n+1):
    for j in range(1, n+1):
        for k in range(1,n+1):
            if (i+j+k == n) and (i!=j or j!=k):
                if sorted(list(set(itertools.permutations([i,j,k],3)))) not in arr:
                    arr.append(sorted(list(set(itertools.permutations([i,j,k],3)))))

print(arr)       
arr1 = [arr[i][j] for i in range(len(arr)) for j in range(len(arr[i]))]
print(arr1)
print(len(arr1))

I solved it on my own even though I am relatively new to advanced programming and I know that my code is not optimized. I would like to optimize it as well as make the list of the different ways in a single line using list comprehension.

My list comprehension has errors:

arr = [sorted(list(set(itertools.permutations([i,j,k],3))) if ((i+j+k == n) and (i!=j or j!=k) and sorted(list(set(itertools.permutations([i,j,k],3))))) not in arr for i in range(1,n+1) for j in range(1, n+1) for k in range(1,n+1)]

Any kind of help will be appreciated!


Solution

  • I would recommend that you avoid too many nested conditions and loops. And is there a reason why you want your answer to involve a list comprehension? I like using them too, but beyond a point they become too long to, um... comprehend. Here's an alternative solution that uses fewer loops, avoids long list comprehensions and is easier to read -- to my eyes, at least.

    In [29]: from itertools import permutations, combinations_with_replacement
    
    In [30]: def all_valid_permutations(candies, members):
        ...:     combos = combinations_with_replacement(range(1, candies), members)
        ...:     valid_permutations = set()
        ...:     for item in combos:
        ...:         for perm in permutations(item):
        ...:             if sum(perm) == candies and len(set(perm)) >= members-1:
        ...:                 valid_permutations.add(perm)
        ...:
        ...:     return valid_permutations
        ...:
    
    In [31]: all_valid_permutations(6, 3)
    Out[31]:
    {(1, 1, 4),
     (1, 2, 3),
     (1, 3, 2),
     (1, 4, 1),
     (2, 1, 3),
     (2, 3, 1),
     (3, 1, 2),
     (3, 2, 1),
     (4, 1, 1)}
    

    If you know your basic combinatorics you can guess what the imported functions do. (Or you can read the docs, if you prefer.) I can't guarantee performance, though, without fully knowing your use-case. This is just a quick solution, off the top of my head. I bet you can optimise this further.