using System.Collections.Generic;
using Microsoft.VisualStudio.TestTools.UnitTesting;
namespace CodingInterview
{
[TestClass]
public class CreateListsOfNodesForEachLevel
{
public List<List<BinaryTreeNode<T>>> GetListOfNodes<T>(BinaryTreeNode<T> root)
{
var result = new List<List<BinaryTreeNode<T>>>();
GetListOfNodesInternal(root, 0, result);
return result;
}
public void GetListOfNodesInternal<T>(BinaryTreeNode<T> root, int level, List<List<BinaryTreeNode<T>>> result)
{
if (root == null)
{
return;
}
if (result.Count <= level)
{
result.Add(new List<BinaryTreeNode<T>>());
}
result[level].Add(root);
GetListOfNodesInternal<T>(root.LeftNode, level + 1, result);
GetListOfNodesInternal<T>(root.RightNode, level + 1, result);
}
#region Unit tests
[TestMethod]
public void GetListOfNodesTest()
{
var root = new BinaryTreeNode<string>("1")
{
LeftNode =
new BinaryTreeNode<string>("2_Left")
{
LeftNode = new BinaryTreeNode<string>("3_Left_1"),
RightNode = new BinaryTreeNode<string>("3_Right_1")
},
RightNode = new BinaryTreeNode<string>("2_Right") {RightNode = new BinaryTreeNode<string>("3_Right_2")}
};
var result = GetListOfNodes(root);
Assert.IsNotNull(result);
Assert.AreEqual(3, result.Count);
Assert.AreEqual(1, result[0].Count);
Assert.AreEqual(2, result[1].Count);
Assert.AreEqual(3, result[2].Count);
}
[TestMethod]
public void GetListOfNodesDeepTreeTest()
{
var root = new BinaryTreeNode<string>("1")
{
LeftNode =
new BinaryTreeNode<string>("2_Left")
{
LeftNode = new BinaryTreeNode<string>("3_Left_1"),
RightNode = new BinaryTreeNode<string>("3_Right_1")
}
};
root.RightNode = new BinaryTreeNode<string>("2_Right")
{
RightNode =
new BinaryTreeNode<string>("3_Right_2")
{
RightNode =
new BinaryTreeNode<string>("4_Right")
{
RightNode =
new BinaryTreeNode<string>("5_Right")
{
RightNode =
new BinaryTreeNode<string>("6_Right")
{
LeftNode = new BinaryTreeNode<string>("7_Left")
}
}
}
}
};
var result = GetListOfNodes(root);
Assert.IsNotNull(result);
Assert.AreEqual(7, result.Count);
Assert.AreEqual(1, result[0].Count);
Assert.AreEqual(2, result[1].Count);
Assert.AreEqual(3, result[2].Count);
Assert.AreEqual(1, result[3].Count);
Assert.AreEqual(1, result[4].Count);
Assert.AreEqual(1, result[5].Count);
Assert.AreEqual(1, result[6].Count);
}
#endregion
}
}