optimizations to IntPairMap.java

[sbp.git] / src / edu / berkeley / sbp / util / HashMapBag.java

diff --git a/src/edu/berkeley/sbp/util/HashMapBag.java b/src/edu/berkeley/sbp/util/HashMapBag.java
index 6f5585d..210290d 100644 (file)
--- a/src/edu/berkeley/sbp/util/HashMapBag.java
+++ b/src/edu/berkeley/sbp/util/HashMapBag.java
@@ -1,3 +1,5 @@
+// Copyright 2006-2007 all rights reserved; see LICENSE file for BSD-style license
+
 package edu.berkeley.sbp.util;
 import java.util.*;
 
@@ -9,7 +11,9 @@ public final class HashMapBag<K,V> implements MapBag<K,V> {
     public void add(K k, V v) {
         HashSet<V> hs = hm.get(k);
         if (hs==null) hm.put(k, hs = new HashSet<V>());
+        size -= hs.size();
         hs.add(v);
+        size += hs.size();
     }
 
     public void addAll(K k, Iterable<V> iv) {
@@ -24,7 +28,15 @@ public final class HashMapBag<K,V> implements MapBag<K,V> {
 
     public void remove(K k, V v) {
         if (hm.get(k)==null) return;
-        hm.get(k).remove(v);
+        HashSet<V> hs = hm.get(k);
+        if (hs==null) return;
+        size -= hs.size();
+        hs.remove(v);
+        size += hs.size();
+    }
+
+    public void removeAll(K k, Iterable<V> iv) {
+        for(V v : iv) remove(k, v);
     }
 
     public void clear() { hm.clear(); }
@@ -33,5 +45,14 @@ public final class HashMapBag<K,V> implements MapBag<K,V> {
         return hm.get(k)!=null && hm.get(k).contains(v);
     }
 
+    public void addAll(HashMapBag<K,V> hmb) {
+        for(K k : hmb) addAll(k, hmb.getAll(k));
+    }
+    public void removeAll(HashMapBag<K,V> hmb) {
+        for(K k : hmb) removeAll(k, hmb.getAll(k));
+    }
+
     public Iterator<K> iterator() { return hm.keySet().iterator(); }
+    public int size() { return size; }
+    private int size = 0;
 }