“咲球问题”Python代码练习

昨晚小柠檬与咲夜酱在Mastodon上闲聊的时候，偶然提出了一道很有趣的题目。该题目抽象一下可以这样表述：

将n个小球放入若干个不同的盒子中，每个盒子里的小球数不少于1且最多不超过k。求一共有多少种放置小球的组合？

看上去与“n个球装进m个盒子里”的经典小球问题挺像，但还是有些不同，为示区别并纪念该问题的源起，暂以“咲球问题”称之。

于是“咲球问题”一出，聊天顿时聊不成了，几个人愣是琢磨了一晚上这题该怎么解。此题初看之下本以为并不太难，但仔细一想发现还是蛮复杂的，光取小数拿纸笔枚举都是一件费劲的事，如果要区分小球的话就更费劲了。不过正好提供了一个机会可以拿来练手新学的Python，不挑战一下还是人，于是在经历无数次失败后，终于使出吃奶的劲，写出下面的代码解决了这个问题：

# 首先考虑不对小球进行区分的情况

import itertools, math

def GetNumList(n, k):
    """
    n, k：正整数，n >= k
    返回：二级嵌套列表，大列表中嵌套若干个小列表，例：[[1,2], [1,1,1]]
    该函数得到一个可能的数目列表，不包含排序
    """
    resultList = []
    # k = 1，生成一个全是1的列表
    if k == 1: 
        l = []
        for i in range(n):
            l += [1]
        resultList = [l]
    # k > 1
    else:
    # GetNumList(n, k) = GetNumList(n, k-1) + 有至少一个数等于k时的组合
        r = n // k # n除以k的整数部分
        for a in range(1, r + 1): # a为k出现的次数，范围为1, 2, ..., r
            recurList = GetNumList(n - a * k, k - 1) # 递归
            # 遍历recuList中的每个小列表
            for l in recurList:
                # 有几个k就在小列表中添加几次k
                for each_a in range(a):
                    l.append(k)
            resultList += recurList # 遍历完recuList后，添加到resultList
        resultList += GetNumList(n, k - 1)
    return resultList

def GetCombList(n, k):
    """
    n, k：正整数，n >= k
    返回：二级嵌套列表，所有数字排列枚举（相同元素看做相同的排列）
    如果不区分小球，该列表长度即为组合数
    """
    resultList = GetNumList(n, k)
    combList = []
    for i in range(len(resultList)):
        for item in itertools.permutations(resultList[i]):
            if list(item) not in combList:
                combList.append(list(item))
    return combList

# 现在考虑区分小球，如，将n个小球标记为A1, A2, ..., An
# 则对于GetCombList(n, k)中的每一个小列表，如[2, 1, 1]，都存在将A1, A2, ..., An分配为列表组合的排列

def CalPermNumPerComb(l, n):
    """
    l：由正整数组成的单层列表；n：正整数，l中所有元素相加等于n
    返回：正整数，该列表所有的排列数
    """
    if n == 0:
        return 1
    else:
        c = math.factorial(n) // (math.factorial(l[0]) * math.factorial(n - l[0]))
        return c * CalPermNumPerComb(l[1:], n - l[0])

# 最终的函数
def CalAllPermNum(n, k):
    """
    n, k：正整数，n >= k
    返回：正整数，所有的排列数
    """
    count = 0
    combList = GetCombList(n, k)
    for l in combList:
        count += CalPermNumPerComb(l, n)
    return count

回头看的话，这段代码写得不严谨，特别是GetNumList(n-a*k, k-1)这里，其实是比较容易出事故的。虽说从题目看，毫无疑问n和k都是正整数且n >= k，但GetNumList(n-a*k, k-1)几乎是必然会出现n比k小，甚至n = 0的情况。我在写的时候根本没有考虑过这个问题，只想了k临界值的情况，万幸这段代码中没有需要n >= k的部分，且在Python中，如果for循环的范围是0，那么并不会报错，只是不会执行这个循环。

当然，作为一个刚学代码的萌新，虽然我已经尽力，但难免写得非常愚蠢。而小柠檬大佬只写了两行就达到了同样的效果：

import functools, math

CalComb = lambda n, k: [ [i] + arr for i in range(1, min(n,k)+1) for arr in CalComb(n-i,k) ] if n != 0 else [[]]

def AllComb(n, k): return sum(functools.reduce((lambda x,y: x//math.factorial(y)), s, math.factorial(n)) for s in CalComb(n, k))

这两行代码真是看得我一愣一愣现在都没缓过劲来，根本想不到还能这样递归，大概这就是萌新和大佬的差距吧……

但不管怎么说，虽然写得比较愚蠢，但好歹也是写出来了，既然写出来就要测试一下。取n为[1,5]，测试输出的结果：

print("n,k,组合数")
for n in range(1, 6):
    for k in range(1, n+1):
        print("{0},{1},{2}".format(n, k, CalAllPermNum(n, k)))

可得到如下结果：

n	k	组合数
1	1	1
2	1	2
2	2	3
3	1	6
3	2	12
3	3	13
4	1	24
4	2	66
4	3	74
4	4	75
5	1	120
5	2	450
5	3	530
5	4	540
5	5	541

输出的时候也能很明显地感觉到我写的代码时间复杂度太高，n为[1,5]的时候尚且没有明显体感，但当n取到10的时候，计算起来就非常非常慢了。而执行小柠檬的代码一瞬间就出了结果，不愧是大佬写代码。

不知道这个结果在数学问题上有没有什么叫法，但当k=n时，会变成这样一组数列：1, 3, 13, 75, 541, ...

该数列被称作“Ordered Bell numbers”或“Fubini numbers”，OEIS上给出了该数列的解法。

据说小柠檬大佬还在兴致勃勃地研究“咲球问题”是如何和这个数列扯上关系的。我对纯粹的数学问题就没太多兴趣了，不过拿一个突发奇想的题目练习一下computational thinking还是很有趣的体验。虽说体验完以后明显地感觉到了我和专业人士之间马里亚纳海沟一般的差距……不过这也是显然的，否则人家专业人士不用混了。希望自己有朝一日也能写出如此流畅的代码。