Make classes enumerable

basp1 · Jul 12, 2019 · bb79bc2 · bb79bc2
1 parent 6c9c50e
commit bb79bc2
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 8 deletions.
diff --git a/src/FlatMap.cs b/src/FlatMap.cs
@@ -1,10 +1,11 @@
 using System;
-using System.Collections.Generic;
 using System.Linq;
+using System.Collections;
+using System.Collections.Generic;
 
 namespace sharpcode
 {
-    public class FlatMap<Key, Value>
+    public class FlatMap<Key, Value> : IEnumerable<KeyValuePair<Key, Value>>
     {
         List<Key> keys;
         List<Value> values;
@@ -160,5 +161,18 @@ int LowerBound(Key key)
 
             return l;
         }
+
+        public IEnumerator<KeyValuePair<Key, Value>> GetEnumerator()
+        {
+            for (int i = 0; i < Count; i++)
+            {
+                yield return new KeyValuePair<Key, Value>(keys[i], values[i]);
+            }
+        }
+
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
     }
 }
diff --git a/src/SparseArray.cs b/src/SparseArray.cs
@@ -1,9 +1,10 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 
 namespace sharpcode
 {
-    public class SparseArray<T>
+    public class SparseArray<T> : IEnumerable<KeyValuePair<int, T>>
     {
         int size;
         T @default;
@@ -20,15 +21,15 @@ public SparseArray(IList<T> init, T @default = default) : this(init.Count, @defa
         {
             for (int i = 0; i < size; i++)
             {
-                Insert(i, init[i]);
+                Set(i, init[i]);
             }
         }
 
         public SparseArray(int size, IEnumerable<KeyValuePair<int, T>> init, T @default = default) : this(size, @default)
         {
             foreach (var item in init)
             {
-                Insert(item.Key, item.Value);
+                Set(item.Key, item.Value);
             }
         }
 
@@ -47,7 +48,7 @@ public bool Equals(SparseArray<T> that)
             return items.Equals(that.items);
         }
 
-        public void Insert(int index, T value)
+        public void Set(int index, T value)
         {
             if (index < 0 || index >= size)
             {
@@ -67,7 +68,7 @@ public bool Contains(int index)
             return items.Contains(index);
         }
 
-        public int Nnz
+        public int Count
         {
             get
             {
@@ -129,8 +130,21 @@ public T this[int index]
 
             set
             {
-                Insert(index, value);
+                Set(index, value);
             }
         }
+
+        public IEnumerator<KeyValuePair<int, T>> GetEnumerator()
+        {
+            foreach (var item in items)
+            {
+                yield return item;
+            }
+        }
+
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
     }
 }