-

cool-RR · cool-RR · commit 5da01d0ecd65 · 2013-09-17T23:58:58.000+03:00
diff --git a/source_py2/python_toolbox/logic_tools.py b/source_py2/python_toolbox/logic_tools.py
@@ -3,6 +3,8 @@
 
 '''This module defines logic-related tools.'''
 
+import collections
+
 from python_toolbox import cute_iter_tools
 from python_toolbox import sequence_tools
 
@@ -30,6 +32,26 @@ def all_equal(iterable, exhaustive=False):
     return all(a==b for (a, b) in pairs)
 
 
+def get_equivalence_classes(iterable, key):
+    '''
+    Divide items in `iterable` to equivalence classes, using the key function.
+    
+    i.e. each item will be put in a set with all other items that had the same
+    result when put through the `key` function.
+    
+    Returns a `dict` with keys being the results of the function, and the
+    values being the sets of items with those values.
+    '''
+    key_function = \
+             (lambda item: getattr(item, key)) if isinstance(key, str) else key
+    equivalence_class_to_members = collections.defaultdict(set)
+    for item in iterable:
+        equivalence_class = key_function(item)
+        equivalence_class_to_members[equivalence_class].add(item)
+            
+    return equivalence_class_to_members
+        
+      
 def logic_max(iterable, relation=lambda a, b: (a >= b)):
     '''
     Get a list of maximums from the iterable.
diff --git a/source_py2/test_python_toolbox/test_logic_tools/test_get_equivalence_classes.py b/source_py2/test_python_toolbox/test_logic_tools/test_get_equivalence_classes.py
@@ -0,0 +1,18 @@
+# Copyright 2009-2013 Ram Rachum.
+# This program is distributed under the MIT license.
+
+import itertools
+
+from python_toolbox.logic_tools import get_equivalence_classes
+
+
+def test():
+    assert get_equivalence_classes([1, 2, 3, 1j, 2j, 3j, 1+1j, 2+2j, 3+3j],
+                                   abs) == {
+        1: {1, 1j},
+        2: {2, 2j},
+        3: {3, 3j},
+        2**0.5: {1 + 1j},
+        8**0.5: {2 + 2j},
+        18**0.5: {3 + 3j},
+    }
