leetcode/main.py at master · woowenjun99/leetcode · 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
37
38
39
40
41
42
43
44
45
from sortedcontainers import SortedSet
from typing import List

class TaskManager:
    def __init__(self, tasks: List[List[int]]):
        self.task_id_to_priority: dict[int, int] = {}
        self.task_id_to_user_id: dict[int, int] = {}
        self.sorted_priority_task_id = SortedSet()
        for user_id, task_id, priority in tasks:
            self.task_id_to_user_id[task_id] = user_id
            self.task_id_to_priority[task_id] = priority
            self.sorted_priority_task_id.add((-priority, -task_id))

    def add(self, userId: int, taskId: int, priority: int) -> None:
        self.task_id_to_user_id[taskId] = userId
        self.task_id_to_priority[taskId] = priority
        self.sorted_priority_task_id.add((-priority, -taskId))

    def edit(self, taskId: int, newPriority: int) -> None:
        current_priority = self.task_id_to_priority[taskId]
        self.sorted_priority_task_id.remove((-current_priority, -taskId))
        self.task_id_to_priority[taskId] = newPriority
        self.sorted_priority_task_id.add((-newPriority, -taskId))

    def rmv(self, taskId: int) -> None:
        current_priority = self.task_id_to_priority[taskId]
        self.sorted_priority_task_id.remove((-current_priority, -taskId))
        del self.task_id_to_priority[taskId]
        del self.task_id_to_user_id[taskId]

    def execTop(self) -> int:
        if not self.sorted_priority_task_id: return -1
        _, task_id = self.sorted_priority_task_id[0]
        user_id = self.task_id_to_user_id[-task_id]
        self.rmv(-task_id)
        return user_id

m = TaskManager([[1, 101, 10], [2, 102, 20], [3, 103, 15]])
print(m.sorted_priority_task_id)
m.add(4, 104, 5)
print(m.sorted_priority_task_id)
m.edit(102, 8)
print(m.sorted_priority_task_id)
m.execTop()
print(m.sorted_priority_task_id)