【JD算法题】定义一个数组的权值为，该数组最大值的出现次数。求长度为n且每个元素范围都在[1,n]的所有数组的权值之和。

Problem

小红定义一个数组的权值为，该数组最大值的出现次数。
例如[2,3,3,4]的权值为1，[2,3,3,3]的权值为3.
小红想知道，长度为n，且每个元素范围都在[1,n]的数组（显然有n^n个数组），这些数组的权值之和是多少？由于该数可能过大，请对$10^9+7$取模。

输入描述：一个正整数n。1<=n<=1000
输出描述：所有数组的权值之和。

示例：输入：2； 输出：6

思考

1. 采用回溯法，像写排列数一样列出所有情况。会超时
2. 考虑更普适的方法，以内存换时间：

考虑长度为n的数组，可能的权值为1,2,3,4，…, n；其中，每种权值可以搭配 “该数组的最大值” 为1,2,3,4，…, n的情况。由此可以构建一个$n\times n$的矩阵。

举个例子，n=5的情况：

权值	max=1	max=2	max=3	max=4	max=5
1	0	$C_5^1\ast M_4^1$	$C_5^1\ast M_4^2$	$C_5^1\ast M_4^3$	$C_5^1\ast M_4^4$
2	0	$C_5^2\ast M_3^1$	$C_5^2\ast M_3^2$	$C_5^2\ast M_3^3$	$C_5^2\ast M_3^4$
3	0	$C_5^3\ast M_2^1$	$C_5^3\ast M_2^2$	$C_5^3\ast M_2^3$	$C_5^3\ast M_2^4$
4	0	$C_5^4\ast M_1^1$	$C_5^4\ast M_1^2$	$C_5^4\ast M_1^3$	$C_5^4\ast M_1^4$
5	1	$C_5^5$	$C_5^5$	$C_5^5$	$C_5^5$

其中， $M_i^j$ 代表 有$i$ 个位置，要填$j$ 个数。即有$M_i^j=j^i$. $C_i^j$ 是组合数。

以权值为2，max=4为例，代表数组中4出现了2次，5没出现，数组剩下的3个位置由‘1,2,3’三个数来填，所以有$C_5^2\ast M_3^3$种情况。

如果直接简单粗暴的按这个思路写代码：

from itertools import combinations
def C(a,b):
    A = [i for i in range(0,a)]
    res = list(combinations(A, b))
    return len(res)

def M(a,b):
    return pow(b, a)

def Problem3(n):
    N_matrix = []
    res = 0
    for i in range(1,n):   # row
        row_i = [0]
        for j in range(1, n):
            row_i.append(C(n, i) * M(n-i, j))
        N_matrix.append(row_i)
        res+=i * sum(row_i)
    N_matrix.append([1]*n)
    res += n*n
    return

虽然可以算出结果，但当n=20时就已经需要2秒了，时间复杂度呈指数式增长。

注意到矩阵的第一行~倒数第二行的M乘数部分呈现固定倍数关系，比如上述矩阵第4行为：
$0$, $C_5^4\ast 1^1$, $C_5^4\ast 2^1$, $C_5^4\ast 3^1$, $C_5^4\ast 4^1$

第三行为：
$0$, $C_5^3\ast 1^2$, $C_5^3\ast 2^2$, $C_5^3\ast 3^2$, $C_5^3\ast 4^2$

不看组合数部分，则将第四行乘上$0，1，2，3，4$可以得到第三行；同理，将第三行乘上$0，1，2，3，4$可以得到第二行。

加速后的代码为：

def Problem3_speed(n):
    multiples = [i for i in range(n)]
    a = [1]*n
    row_n_i = a
    res = n * n
    for i in range(1, n):
        row_n_i = cdot(row_n_i, multiples)
        coeff = C(n,n-i)*(n-i)
        res += coeff*sum(row_n_i)
    return res

发现算n=20，已经可以由第一版本的代码的2秒加速成0.1秒，但由于组合数计算涉及大量阶乘依然很费时间。组合数公式：
$$C_m^n=\frac{m!}{n!(m-n)!}$$
如果提前把$[1!,2!,3!,...,n!]$存下来，就避免了大量的重复计算。计算$[1!,2!,3!,...,n!]$的代码：

# [1!, 2!, 3!, ..., n!]
def factorial(n):
    a=1
    n_fac_list = []
    #for循环遍历
    for i in range(1,n+1):
        a*=i
        n_fac_list.append(a)
    return n_fac_list

而且
$$C_m^n=C_m^{m-n}$$
注意到每一列的组合数都是一致的，下面写个函数，提前将这些组合数（Cn1, Cn2,Cn3,…, Cnn）存下来：

def C_table(n):
    '''
    :param n:
    :return: from row1 to row n: Cn1, Cn2,Cn3,..., Cnn
    '''
    c_table = []
    n_fac_table = factorial(n)
    if n % 2 == 0:
        for i in range(1, n//2+1):
            curr_num = n_fac_table[-1] // (n_fac_table[i-1] * n_fac_table[n-i-1])
            c_table.append(curr_num)
        c_table+=c_table[::-1][1:]
        c_table+=[1]
    else:
        for i in range(1, n // 2 + 1):
            curr_num = n_fac_table[-1] // (n_fac_table[i - 1] * n_fac_table[n - i - 1])
            c_table.append(curr_num)
        c_table += c_table[::-1]
        c_table += [1]
    return c_table

下面再写主函数的代码：

def Problem3_speed2(n):
    mods = pow(10,9)+7
    multiples = [i for i in range(n)]
    a = [1]*n
    c = C_table(n)
    row_n_i = a
    res = n * n
    for i in range(1, n):
        row_n_i = cdot(row_n_i, multiples)
        coeff = (n-i) * c[n-i-1]
        res += coeff*sum(row_n_i)
    return res%mods

测试一下时间：

if __name__ =='__main__':
    import time
    t2 = time.time()
    print(Problem3_speed2(1000))
    t3 = time.time()
    print(t3-t2)

n=1000也只需要0.4秒 ^_^