Pass TestFiniteMakeSetDisjointUnion

1 files changed, 42 insertions, 2 deletions

diff --git a/src/act4e_solutions/sets_sum.py b/src/act4e_solutions/sets_sum.py
index eab93d3..9a3f364 100644
--- a/src/act4e_solutions/sets_sum.py
+++ b/src/act4e_solutions/sets_sum.py
@@ -1,10 +1,50 @@
-from typing import Any, Sequence, TypeVar
+from typing import Any, Sequence, TypeVar, List, Tuple
 
 import act4e_interfaces as I
 
 X = TypeVar("X")
+C = TypeVar("C")
+E = Tuple[int, C]
+
+class MyFiniteSetDisjointUnion(I.FiniteSetDisjointUnion[C, E]):
+    _components: List[I.FiniteSet[C]]
+    def __init__(self, components: Sequence[I.FiniteSet[C]]):
+        self._components = list(components)
+
+    def contains(self, e: E):
+        i, val = e
+        return self._components[i].contains(val)
+
+    def size(self):
+        return sum(c.size() for c in self._components)
+
+    def save(self, h: I.IOHelper, x: E) -> List[E]:
+        return cast(I.ConcreteRepr, list(x))
+
+    def load(self, h: I.IOHelper, o: I.ConcreteRepr):
+        return cast(E, o)
+
+    def elements(self) -> Sequence[E]:
+        for i, c in enumerate(self._components):
+            for e in c.elements():
+                yield [i, c]
+
+    def components(self) -> Sequence[I.FiniteSet[C]]:
+        return self._components
+
+    def pack(self, i: int, e: C) -> E:
+        if i < 0 or i > len(self._components):
+            raise I.InvalidValue
+
+        # if not self._components[i].contains(e):
+        #     raise I.InvalidValue
+
+        return [i, e]
+
+    def unpack(self, e: E) -> Tuple[int, C]:
+        return tuple(e)
 
 
 class SolFiniteMakeSetDisjointUnion(I.FiniteMakeSetDisjointUnion):
     def disjoint_union(self, components: Sequence[I.FiniteSet[X]]) -> I.FiniteSetDisjointUnion[X, Any]:
-        raise NotImplementedError()  # implement here
+        return MyFiniteSetDisjointUnion(components)