用于查找不同元素集的 Python 时间优化

Question

我有一个看起来像这样的Python硬件：

找出每件衣服有不同颜色的衣服套装的数量。

输入 = n（服装数量）

对于接下来的 n 行，输入将是衣服的编号（帽子 = 1，T 恤 = 2，裤子 = 3），下一个数字是颜色。

预期输出为：不同服装的数量。

这是所需输入/输出的示例：

输入：

9
1 1
2 1
3 1
1 5
2 5
3 5
1 9
2 9
3 9

输出 =

6

我编写了此处所示的代码，并且它有效！

但是当 n 大于 100000 时，会出现超出限制时间的错误。时间限制为1s，内存限制为256 mb。

import itertools
n = int(input())
pants, hats, t_shirts = [], [], []
for i in range(n):
    a, b = map(int, input().split())
    if(a==1):
        hats.append(b)
    elif(a==2):
        t_shirts.append(b)
    elif(a==3):
        pants.append(b)

result = 0

for i in itertools.product(pants, hats, t_shirts):
    if len(set(i)) == 3:
        result += 1
print(result)

面对错误后，我决定找一个像这样的硬件示例来进行测试，并且不需要手动输入数据。所以，我编写了这个代码片段并尝试优化求解时间。

import time
import itertools

hats = [n for n in range(200)]
cats = [n for n in range(100)]
bats = [n for n in range(100)]

result = 0
starttime = time.time()

for i in itertools.product(hats, bats, cats):
    if len(set(i)) == 3:
        result += 1
print(result)
print(time.time() - starttime)

当前输出为：

1960200
1.357010841369629

有人可以帮忙将求解时间优化到 1 秒以下吗？也允许使用所有其他编程语言。谢谢

Answer 1

也允许使用所有其他编程语言。

那么 C 版本应该优于 Python 代码片段。

#include <stdio.h>

int main()
{
    int hats[200], cats[100], bats[100];
    for (int i = 0; i < 200; ++i) hats[i] = i;
    for (int i = 0; i < 100; ++i) cats[i] = i;
    for (int i = 0; i < 100; ++i) bats[i] = i;
    int result = 0;
    for (int h = 0; h < 200; ++h)
    for (int c = 0; c < 100; ++c)
    for (int b = 0; b < 100; ++b)
        result += hats[h]!=cats[c] && hats[h]!=bats[b] && cats[c]!=bats[b];
    printf("%d\n", result);
}

但即使如此，对于 n 大于 100000 来说也可能需要大约一天的时间。

所以我建议根据每种布料类型的独特颜色的数量和类型之间的共同颜色的数量来计算结果。这段 Python 代码似乎是这样做的 - 它乘以不同颜色的数量并减去相同颜色的组合数量：

H = set(hats)
C = set(cats)
B = set(bats)
h = len(H)
c = len(C)
b = len(B)
result = (h*c-len(H&C))*b-c*len(H&B)-h*len(C&B)+2*len(H&C&B)