add merge_coint

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Add `jax_interface`
 
+- Add `merge_count` in `results` module
+
 ## 1.1.0
 
 ### Added

tensorcircuit/results/counts.py

@@ -10,8 +10,6 @@
 Tensor = Any
 ct = Dict[str, int]
 
-# TODO(@refraction-ray): merge_count
-
 
 def reverse_count(count: ct) -> ct:
     ncount = {}
@@ -107,6 +105,22 @@ def expectation(
     return r / shots
 
 
+def merge_count(*counts: ct) -> ct:
+    """
+    Merge multiple count dictionaries by summing up their counts
+
+    :param counts: Variable number of count dictionaries
+    :type counts: ct
+    :return: Merged count dictionary
+    :rtype: ct
+    """
+    merged: ct = {}
+    for count in counts:
+        for k, v in count.items():
+            merged[k] = merged.get(k, 0) + v
+    return merged
+
+
 def plot_histogram(data: Any, **kws: Any) -> Any:
     """
     See ``qiskit.visualization.plot_histogram``:

tests/test_results.py

@@ -15,6 +15,39 @@ def test_marginal_count():
     assert counts.marginal_count(d, [2, 1, 0])["001"] == 4
 
 
+def test_merge_count():
+
+    c1 = {"00": 10, "01": 20, "11": 30}
+    c2 = {"00": 5, "10": 15, "11": 25}
+    c3 = {"01": 10, "10": 20}
+
+    # Test merging two count dicts
+    merged = counts.merge_count(c1, c2)
+    assert merged["00"] == 15
+    assert merged["01"] == 20
+    assert merged["10"] == 15
+    assert merged["11"] == 55
+
+    # Test merging three count dicts
+    merged = counts.merge_count(c1, c2, c3)
+    assert merged["00"] == 15
+    assert merged["01"] == 30
+    assert merged["10"] == 35
+    assert merged["11"] == 55
+
+    # Test merging single count dict
+    merged = counts.merge_count(c1)
+    assert merged == c1
+
+    # Test merging empty dicts
+    merged = counts.merge_count({}, {})
+    assert merged == {}
+
+    # Test merging empty with non-empty
+    merged = counts.merge_count({}, c1)
+    assert merged == c1
+
+
 def test_count2vec():
     assert counts.vec2count(counts.count2vec(d, normalization=False), prune=True) == d