Made KdTree Enumerable

Author: codeandcats

KdTreeLib/IKdTree.cs

@@ -6,7 +6,7 @@
 
 namespace KdTree
 {
-	public interface IKdTree<TKey, TValue>
+	public interface IKdTree<TKey, TValue> : IEnumerable<KdTreeNode<TKey, TValue>>
 	{
 		bool Add(TKey[] point, TValue value);
 

KdTreeLib/KdTree.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -553,5 +554,66 @@ public static KdTree<TKey, TValue> LoadFromFile(string filename)
 			}
 
 		}
-	}
+
+        public IEnumerator<KdTreeNode<TKey, TValue>> GetEnumerator()
+        {
+            var left = new Stack<KdTreeNode<TKey, TValue>>();
+            var right = new Stack<KdTreeNode<TKey, TValue>>();
+
+            Action<KdTreeNode<TKey, TValue>> addLeft = node =>
+            {
+                if (node.LeftChild != null)
+                {
+                    left.Push(node.LeftChild);
+                }
+            };
+
+            Action<KdTreeNode<TKey, TValue>> addRight = node =>
+            {
+                if (node.RightChild != null)
+                {
+                    right.Push(node.RightChild);
+                }
+            };
+
+            if (root != null)
+            {
+                yield return root;
+
+                addLeft(root);
+                addRight(root);
+
+                while (true)
+                {
+                    if (left.Any())
+                    {
+                        var item = left.Pop();
+
+                        addLeft(item);
+                        addRight(item);
+
+                        yield return item;
+                    }
+                    else if (right.Any())
+                    {
+                        var item = right.Pop();
+
+                        addLeft(item);
+                        addRight(item);
+
+                        yield return item;
+                    }
+                    else
+                    {
+                        break;
+                    }
+                }
+            }
+        }
+
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+    }
 }

KdTreeTestsLib/KdTreeTests.cs

@@ -2,7 +2,6 @@
 using System.Collections.Generic;
 using System.Linq;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
-using KdTree;
 using KdTree.Math;
 
 struct City
@@ -266,5 +265,23 @@ public void TestGetNearestNeighbours()
 				}
 			}
 		}
-	}
+
+        [TestMethod]
+        [TestCategory("KdTree")]
+        public void TestEnumerable()
+        {
+            AddTestNodes();
+
+            foreach (var node in tree)
+            {
+                var testNode = testNodes.FirstOrDefault(n => n.Point == node.Point && n.Value == node.Value);
+
+                Assert.IsNotNull(testNode);
+
+                testNodes.Remove(testNode);
+            }
+
+            Assert.AreEqual(0, testNodes.Count);
+        }
+    }
 }