'How to improve efficiency of appending to and then checking for subsequent elements in list

i'm attempting to solve: http://orac.amt.edu.au/cgi-bin/train/problem.pl?problemid=980&set=aio17int

The algorithm exceeds the time limit for larger files. I'm fairly new to python, and can't find a more efficient solution online for checking if an element is in a list and then appending accordingly. Please help improve the speed without using any imports. Thank you.

input file:

8 5
2 7
1 8
8 4
7 5
8 6
re = 1
bl = 1
red = [1]
blue = [2]
input_file = open("tagin.txt","r")
n, m = map(int, input_file.readline().split())
for i in range(m):
    a, b = map(int, input_file.readline().split())
    if a in red:
        red.append(b)
        re+=1
    elif a in blue:
        blue.append(b)
        bl+=1

output_file = open("tagout.txt", "w")
output_file.write(str(re) + " " + str(bl))
output_file.close()

output file:

4 3

Also please advise if stack overflow is the wrong platform to ask this question, if so what should I use?



Solution 1:[1]

If I understand the problem correctly then this should fulfil the objective:

with open('tagin.txt') as tagin:
    _, M = map(int, next(tagin).split())
    red = {1}
    blue = {2}
    for _ in range(M):
        a, b = map(int, next(tagin).split())
        (red if a in red else blue).add(b)
    print(len(red), len(blue))

Output:

4 3

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1