algorithms-training/solutions/71.py at master · cuongpiger/algorithms-training · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        colors = dict()
        balls = dict()
        res = []

        for q in queries:
            # print(f"Query: {q}")
            ball, color = q[0], q[1]
            found_ball_color = balls.get(ball, None)

            if found_ball_color is None:
                balls[ball] = color
                set_balls = colors.get(color, set())
                set_balls.add(ball)
                colors[color] = set_balls
            else:
                if found_ball_color != color:
                    set_balls = colors.get(found_ball_color, set())
                    set_balls.discard(ball)
                    colors[found_ball_color] = set_balls

                    if len(set_balls) < 1:
                        colors.pop(found_ball_color, None)


                    set_balls = colors.get(color, set())
                    set_balls.add(ball)
                    colors[color] = set_balls
                    balls[ball] = color
                    # print(f"Color set: {colors}")

            res.append(len(colors))


        return res