Itertools.combinations() 提高超过时间限制

Question

我正在尝试解决这个 Python 编码问题：

描述：安布罗休老师正在教他的二进制数 女儿 Ambrosinita，所以他决定制作一款游戏。

安布罗西奥会给安布罗西尼塔N个号码，安布罗西尼塔可以选择K个号码 在最初的 N 中。对于每个选择的数字，她都会获得一个分数 相当于使用二进制表示的 1 的数量 号码。

帮助 Ambrosinita 找出她可以赚取多少积分。

条目：

第一个输入行包含一个整数T，表示数字 测试用例。

每个测试用例都以包含整数 N（总数）的行开头 和 K（可以选择的数字）。

每个案例的最后输入行包含 N 个整数，表示 Ambrosinita 可以选择的数字。

输出：

为每个测试用例打印一行包含多少点 Ambrosinita 可以赚到。

1≤T≤10

1≤N≤10^3

0≤K≤N

0 ≤ 数字 ≤10^5

执行时间限制为 2 秒。我收到 TLE 错误，但输出与预期相同。因此，输入必须有一定的长度。

这是我的代码：

import itertools

test_cases = int(input())

def binary(num):
    return format(num,'b')

def filter_string_1s(string):
    aux = ''
    for i in string:
        if i == '1':
            aux += i
    return aux

for i in range(test_cases):
    k = input().split()
    k = int(k[1])
    values = input().split()
    values = [int(i) for i in values]
    values = [filter_string_1s(str(binary(i))) for i in values]
    bin_values = []
    for combination in itertools.combinations(values,k):
        aux = 0
        for bin_number in combination:
            for i in bin_number:
                if i == '1':
                    aux += 1
        bin_values.append(aux)
    print(max(bin_values))

问题：我应该采取哪些步骤来优化它以便在执行时限内解决问题？

TLE = 超过时间限制

Answer 1

这是一个“top-k”问题，而不是需要“k”个组合的问题。为了便于阅读，我们选择 N = 990，K = 7。您需要在 990 的列表中选择前 7 个分数。

您通过生成 C(990, 7) = 912,459,983,564,271,542,400 个组合来实现此目的，然后然后确定每个组合中的 7 个数字的每个中有多少 1 位。这就是您没有时间的原因：您有 990 个数字要考虑，但您要为输入中的每个数字重复数万亿次的位数。

把它关掉。您所需要的只是遍历该列表一次 并保持前 7 位计数。列表中的 7 个数字之间没有交互作用。事实上，您甚至不必报告数字，只需报告总分即可。

从七个零的列表开始。现在遍历所有 990 个数字。每当您发现比列表中最小元素大的位计数时（保持排序以方便参考），然后将其替换为新分数（并重新排序）。在所有 990 个号码的末尾，sum 列表。

另外，计算位比您做的要容易得多。将int 转换为二进制字符串并使用str.count(1) 查看其中有多少1 位。

Answer 2

她可以通过选择具有最多 1 位的 K 个数字来获得最大分数。所以你只需要计算每个数字的点数并取前 K。

例如：

import random

def maxPoints(K,N=None,numbers=None):
    numbers = numbers or [random.randrange(0,100001) for _ in range(N)]
    points  = sorted(bin(n).count("1") for n in numbers)
    return sum(points[-K:])

mp = maxPoints(3,numbers=[1,2,3,4,5])
print(mp) # 5

mp = maxPoints(1000,1000) 
print(mp) # will return instantly

Answer 3

解决了！谢谢大家。

代码：

test_cases = int(input())

for i in range(test_cases):
    k = input().split()
    k = int(k[1])
    nums = input().split()
    nums = [format(int(i),'b').count('1') for i in nums]
    nums.sort()
    sum = 0
    for i in range(k):
        sum += nums[-1]
        del nums[-1]
    print(sum)