add keep_counts_lt, keep_counts_gt

src/icepool/expression/multiset_expression.py

@@ -526,6 +526,19 @@ def keep_counts_le(self, n: int, /) -> 'MultisetExpression[T_contra]':
         """
         return icepool.expression.KeepCountsExpression(self, n, operator.le)
 
+    def keep_counts_lt(self, n: int, /) -> 'MultisetExpression[T_contra]':
+        """Keeps counts that are < n, treating the rest as zero.
+
+        For example, `expression.keep_counts_lt(2)` would remove doubles,
+        triplets...
+
+        Example:
+        ```
+        Pool([1, 2, 2, 3, 3, 3]).keep_counts_lt(2) -> [1]
+        ```
+        """
+        return icepool.expression.KeepCountsExpression(self, n, operator.lt)
+
     def keep_counts_ge(self, n: int, /) -> 'MultisetExpression[T_contra]':
         """Keeps counts that are >= n, treating the rest as zero.
 
@@ -542,6 +555,19 @@ def keep_counts_ge(self, n: int, /) -> 'MultisetExpression[T_contra]':
         """
         return icepool.expression.KeepCountsExpression(self, n, operator.ge)
 
+    def keep_counts_gt(self, n: int, /) -> 'MultisetExpression[T_contra]':
+        """Keeps counts that are < n, treating the rest as zero.
+
+        For example, `expression.keep_counts_gt(2)` would remove singles and
+        doubles.
+
+        Example:
+        ```
+        Pool([1, 2, 2, 3, 3, 3]).keep_counts_gt(2) -> [3, 3, 3]
+        ```
+        """
+        return icepool.expression.KeepCountsExpression(self, n, operator.gt)
+
     def keep_counts_eq(self, n: int, /) -> 'MultisetExpression[T_contra]':
         """Keeps counts that are == n, treating the rest as zero.
 

tests/generator_operator_test.py

@@ -69,15 +69,27 @@ def test_example_divide_counts():
     assert result == expected
 
 
+def test_example_keep_counts_le():
+    result = Pool([1, 2, 2, 3, 3, 3]).keep_counts_le(2).expand(Order.Ascending)
+    expected = Die([(1, 2, 2)])
+    assert result == expected
+
+
+def test_example_keep_counts_lt():
+    result = Pool([1, 2, 2, 3, 3, 3]).keep_counts_lt(2).expand(Order.Ascending)
+    expected = Die([(1, )])
+    assert result == expected
+
+
 def test_example_keep_counts_ge():
     result = Pool([1, 2, 2, 3, 3, 3]).keep_counts_ge(2).expand(Order.Ascending)
     expected = Die([(2, 2, 3, 3, 3)])
     assert result == expected
 
 
-def test_example_keep_counts_le():
-    result = Pool([1, 2, 2, 3, 3, 3]).keep_counts_le(2).expand(Order.Ascending)
-    expected = Die([(1, 2, 2)])
+def test_example_keep_counts_gt():
+    result = Pool([1, 2, 2, 3, 3, 3]).keep_counts_gt(2).expand(Order.Ascending)
+    expected = Die([(3, 3, 3)])
     assert result == expected