# 3508. Implement Router
import bisect
from collections import defaultdict, deque
"""
design a data structure that can efficiently manage data packets in a network
router. each data packet consists of the following attributes: 'source',
'destination', 'timestamp'. note that the queries for 'routerAddPacket' will
be made in increasing order of timestamp
"""
class Router(object):
def __init__(self, memoryLimit):
"""
:type memoryLimit: int
"""
self.size = memoryLimit
self.packets = {}
self.counts = defaultdict(list)
self.queue = deque()
def addPacket(self, source, destination, timestamp):
"""
:type source: int
:type destination: int
:type timestamp: int
:rtype: bool
"""
key = self._encode(source, destination, timestamp)
if key in self.packets:
return False
if len(self.packets) >= self.size:
self.forwardPacket()
self.packets[key] = [source, destination, timestamp]
self.queue.append(key)
self.counts[destination].append(timestamp)
return True
def forwardPacket(self):
"""
:rtype: List[int]
"""
if not self.packets:
return []
key = self.queue.popleft()
packet = self.packets.pop(key)
dest = packet[1]
self.counts[dest].pop(0)
return packet
def getCount(self, destination, startTime, endTime):
"""
:type destination: int
:type startTime: int
:type endTime: int
:rtype: int
"""
timestamps = self.counts.get(destination, [])
if not timestamps:
return 0
left = bisect.bisect_left(timestamps, startTime)
right = bisect.bisect_right(timestamps, endTime)
return right - left
def _encode(self, source, destination, timestamp):
return (source << 40) | (destination << 20) | timestamp
if __name__ == "__main__":
obj = Router(3)
print(obj.addPacket(1, 4, 90))
print(obj.addPacket(2, 5, 90))
print(obj.addPacket(1, 4, 90))
print(obj.addPacket(3, 5, 95))
print(obj.addPacket(4, 5, 105))
print(obj.forwardPacket())
print(obj.addPacket(5, 2, 110))
print(obj.getCount(5, 100, 110))
