sumolib.datastructures.OrderedMultiSet

  1# Eclipse SUMO, Simulation of Urban MObility; see https://eclipse.dev/sumo
  2# Copyright (C) 2011-2026 German Aerospace Center (DLR) and others.
  3# This program and the accompanying materials are made available under the
  4# terms of the Eclipse Public License 2.0 which is available at
  5# https://www.eclipse.org/legal/epl-2.0/
  6# This Source Code may also be made available under the following Secondary
  7# Licenses when the conditions for such availability set forth in the Eclipse
  8# Public License 2.0 are satisfied: GNU General Public License, version 2
  9# or later which is available at
 10# https://www.gnu.org/licenses/old-licenses/gpl-2.0-standalone.html
 11# SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-or-later
 12
 13# @file    OrderedMultiSet.py
 14# @author  Jakob Erdmann
 15# @author  Michael Behrisch
 16# @date    2011-10-04
 17
 18"""
 19multi set with insertion-order iteration
 20based on OrderedSet by Raymond Hettinger (c) , MIT-License
 21[http://code.activestate.com/recipes/576694/]
 22"""
 23from __future__ import absolute_import
 24
 25import collections
 26try:
 27    from collections.abc import MutableSet
 28except ImportError:
 29    from collections import MutableSet
 30KEY, PREV, NEXT = range(3)
 31
 32
 33class OrderedMultiSet(MutableSet):
 34
 35    def __init__(self, iterable=None):
 36        self.end = end = []
 37        # sentinel node for doubly linked list
 38        end += [None, end, end]
 39        # key --> [(key, prev1, next1), (key, prev2, next2), ...]
 40        self.map = collections.defaultdict(collections.deque)
 41        self.size = 0
 42        if iterable is not None:
 43            self |= iterable
 44
 45    def __len__(self):
 46        return self.size
 47
 48    def __contains__(self, key):
 49        return key in self.map
 50
 51    def add(self, key):
 52        self.size += 1
 53        end = self.end
 54        curr = end[PREV]
 55        new = [key, curr, end]
 56        curr[NEXT] = end[PREV] = new
 57        self.map[key].append(new)
 58
 59    def discard(self, key):
 60        if key in self.map:
 61            self.size -= 1
 62            deque = self.map[key]
 63            key, prev, next = deque.popleft()
 64            prev[NEXT] = next
 65            next[PREV] = prev
 66            if len(deque) == 0:
 67                self.map.pop(key)
 68
 69    def __iter__(self):
 70        end = self.end
 71        curr = end[NEXT]
 72        while curr is not end:
 73            yield curr[KEY]
 74            curr = curr[NEXT]
 75
 76    def __reversed__(self):
 77        end = self.end
 78        curr = end[PREV]
 79        while curr is not end:
 80            yield curr[KEY]
 81            curr = curr[PREV]
 82
 83    def pop(self, last=True):
 84        if not self:
 85            raise KeyError('set is empty')
 86        key = next(reversed(self)) if last else next(iter(self))
 87        self.discard(key)
 88        return key
 89
 90    def __repr__(self):
 91        if not self:
 92            return '%s()' % (self.__class__.__name__,)
 93        return '%s(%r)' % (self.__class__.__name__, list(self))
 94
 95    def __eq__(self, other):
 96        if isinstance(other, self.__class__):
 97            return len(self) == len(other) and list(self) == list(other)
 98        return set(self) == set(other)
 99
100    def __del__(self):
101        self.clear()                    # remove circular references
102
103    def __sub__(self, other):
104        result = self.__class__()
105        for x in self:
106            result.add(x)
107        for x in other:
108            result.discard(x)
109        return result
110
111    def __or__(self, other):
112        result = self.__class__()
113        for x in self:
114            result.add(x)
115        for x in other:
116            result.add(x)
117        return result