Inline and simplify Remove

BlueRaja · Jun 11, 2017 · b8fc5ee · b8fc5ee
1 parent 570439c
commit b8fc5ee
Showing 1 changed file with 22 additions and 4 deletions.
diff --git a/Priority Queue/SimplePriorityQueue.cs b/Priority Queue/SimplePriorityQueue.cs
@@ -288,13 +288,31 @@ public void Remove(TItem item)
         {
             lock(_queue)
             {
-                SimpleNode removeMe = GetExistingNode(item);
-                if (removeMe == null)
+                SimpleNode removeMe;
+                IList<SimpleNode> nodes;
+                if (item == null)
+                {
+                    if (_nullNodesCache.Count == 0)
+                    {
+                        throw new InvalidOperationException("Cannot call Remove() on a node which is not enqueued: " + item);
+                    }
+                    removeMe = _nullNodesCache[0];
+                    nodes = _nullNodesCache;
+                }
+                else
                 {
-                    throw new InvalidOperationException("Cannot call Remove() on a node which is not enqueued: " + item);
+                    if (!_itemToNodesCache.TryGetValue(item, out nodes) || nodes.Count == 0)
+                    {
+                        throw new InvalidOperationException("Cannot call Remove() on a node which is not enqueued: " + item);
+                    }
+                    removeMe = nodes[0];
+                    if (nodes.Count == 0)
+                    {
+                        _itemToNodesCache.Remove(item);
+                    }
                 }
                 _queue.Remove(removeMe);
-                RemoveFromNodeCache(removeMe);
+                nodes.Remove(removeMe);
             }
         }