python - 问题的时间复杂度和优化方法

Question

我正在解决这个问题problem以贪婪的方式。这是我为解决问题而编写的代码。

import math
n,m = map(int, input().split())

a = list(map(int, input().split()))

a_s = sorted(a,reverse=True)

res = []
for i in range(1,n+1):
    temp = i
    d = 1
    s = 0
    mid = a_s[-temp:]
    l = 0
    while(l<len(mid)):
        s += d*mid[l]
        l += 1
        if  l%m ==0 and l!=0:
            d += 1
    res.append(s)

print(*res,sep=' ')

这里我对 k 的每个值使用循环来计算该值

我正在为此解决方案获得 TLE。谁能帮我以优化的方式解决这个问题？我知道前缀和很有用，但我不知道如何使用它。

感谢您的宝贵时间

Answer 1

Python sum() 对整个列表求和。这对于以下内容很有用。

#!/usr/bin/python
n, m, q, a = 9, 3, 5, sorted([3,14,34,14,7,2,66,7,4], reverse=True)
print(*[sum([v*i+m+q for i, v in enumerate(a[-k:])]) for k in range(1, n+1)])

./eek.py 
8 18 31 48 72 103 148 207 300

只需将公式 i+m+q 替换为其他内容，就可以解决C 问题。吃甜食;)。