Add subtract() method to collections.Counter() objects.

diff --git a/Doc/library/collections.rst b/Doc/library/collections.rst
index a3bbc63cbbedd15509064e122f58ab947eb732b7..5a6f0de9dfa0e5e144b3a013829ec9946d32196b 100644 (file)
--- a/Doc/library/collections.rst
+++ b/Doc/library/collections.rst
@@ -230,6 +230,17 @@ For example::
             >>> Counter('abracadabra').most_common(3)
             [('a', 5), ('r', 2), ('b', 2)]
 
+   .. method:: subtract([iterable-or-mapping])
+
+      Elements are subtracted from an *iterable* or from another *mapping*
+      (or counter).  Like :meth:`dict.update` but subtracts counts instead
+      of replacing them.  Both inputs and outputs may be zero or negative.
+
+            >>> c = Counter(a=4, b=2, c=0, d=-2)
+            >>> d = Counter(a=1, b=2, c=3, d=4)
+            >>> c.subtract(d)
+            Counter({'a': 3, 'b': 0, 'c': -3, 'd': -6})
+
    The usual dictionary methods are available for :class:`Counter` objects
    except for two which work differently for counters.
 

diff --git a/Lib/collections.py b/Lib/collections.py
index 03226696cf1033cafcc92202091d5d9b9975f7a1..9edc294ef0e69122f41172877cbf47b789bd6fe2 100644 (file)
--- a/Lib/collections.py
+++ b/Lib/collections.py
@@ -435,6 +435,34 @@ class Counter(dict):
         if kwds:
             self.update(kwds)
 
+    def subtract(self, iterable=None, **kwds):
+        '''Like dict.update() but subtracts counts instead of replacing them.
+        Counts can be reduced below zero.  Both the inputs and outputs are
+        allowed to contain zero and negative counts.
+
+        Source can be an iterable, a dictionary, or another Counter instance.
+
+        >>> c = Counter('which')
+        >>> c.subtract('witch')             # subtract elements from another iterable
+        >>> c.subtract(Counter('watch'))    # subtract elements from another counter
+        >>> c['h']                          # 2 in which, minus 1 in witch, minus 1 in watch
+        0
+        >>> c['w']                          # 1 in which, minus 1 in witch, minus 1 in watch
+        -1
+
+        '''
+        if iterable is not None:
+            if isinstance(iterable, Mapping):
+                self_get = self.get
+                for elem, count in iterable.items():
+                    self[elem] = self_get(elem, 0) - count
+            else:
+                self_get = self.get
+                for elem in iterable:
+                    self[elem] = self_get(elem, 0) - 1
+        if kwds:
+            self.subtract(kwds)
+
     def copy(self):
         'Like dict.copy() but returns a Counter instance instead of a dict.'
         return Counter(self)

diff --git a/Lib/test/test_collections.py b/Lib/test/test_collections.py
index a7be8d5128f5c47213d946990a85769df35b6260..43305836b5d79461aa1385f8df3139efee00061c 100644 (file)
--- a/Lib/test/test_collections.py
+++ b/Lib/test/test_collections.py
@@ -669,6 +669,17 @@ class TestCounter(unittest.TestCase):
                 set_result = setop(set(p.elements()), set(q.elements()))
                 self.assertEqual(counter_result, dict.fromkeys(set_result, 1))
 
+    def test_subtract(self):
+        c = Counter(a=-5, b=0, c=5, d=10, e=15,g=40)
+        c.subtract(a=1, b=2, c=-3, d=10, e=20, f=30, h=-50)
+        self.assertEqual(c, Counter(a=-6, b=-2, c=8, d=0, e=-5, f=-30, g=40, h=50))
+        c = Counter(a=-5, b=0, c=5, d=10, e=15,g=40)
+        c.subtract(Counter(a=1, b=2, c=-3, d=10, e=20, f=30, h=-50))
+        self.assertEqual(c, Counter(a=-6, b=-2, c=8, d=0, e=-5, f=-30, g=40, h=50))
+        c = Counter('aaabbcd')
+        c.subtract('aaaabbcce')
+        self.assertEqual(c, Counter(a=-1, b=0, c=-1, d=1, e=-1))
+
 class TestOrderedDict(unittest.TestCase):
 
     def test_init(self):

diff --git a/Misc/NEWS b/Misc/NEWS
index 29d9d401d127415ffbde3a4ccb4d5c8d28fd6c06..a557f41b817d9ba113b2bc89f1cc97dba2c3a45c 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -47,6 +47,8 @@ Core and Builtins
 Library
 -------
 
+- collections.Counter() now supports a subtract() method.
+
 - Issue #8294: The Fraction constructor now accepts Decimal and float
   instances directly.