Python time optimization for finding distinct sets of elements

Question

I've got a python HW that looks like this:

find the number of clothes sets that each clothing has a different color.

The input = n (number of costumes)

For the next n line the input will be number of clothe (hat = 1, t_shirt = 2, pants = 3) and the next digit is the color.

The expected output will be: the number of distinct sets of clothing.

Here is a sample of desired input/output:

Input:

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

Output = 6

I wrote this code shown here, and it works!

But for n bigger than 100000 a limit time exceeded error occurs. The time limit is 1s and the memory limit is 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)

After facing the error, I decided to find an example like this HW to be testable and doesn't need to input data manually. so, I wrote this code snippet and tried to optimize solving time.

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)

The current output is :

1960200
1.357010841369629

Can anybody please help to optimize solving time to under 1s? Using all other programming language is also allowed. thanks

Answer 1

using all other programming language is also allowed.

Well, then a C version should outperform the Python code snippet.

#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);
}

But even that may take about a day for n bigger than 100000.

So I suggest to compute the result based on the numbers of unique colors of each cloth type and the numbers of common colors between the types. This Python code seems to do - it multiplies the numbers of distinct colors and subtracts the numbers of combinations with same colors:

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)