Added method to enumerate all suffixes in the suffix tree

Author: Lexey

Experimental/src/CodeJam.Experimental.csproj

@@ -66,6 +66,7 @@
     <Compile Include="Collections\DisjointSets.cs" />
     <Compile Include="Collections\DisjointSetsBase.cs" />
     <Compile Include="Collections\DisjointSetsT.cs" />
+    <Compile Include="Collections\Suffix.cs" />
     <Compile Include="Collections\SuffixTreeBase.cs" />
     <Compile Include="Collections\SuffixTree.cs" />
     <Compile Include="Parsing\BinaryOperator.cs" />

Experimental/src/Collections/Suffix.cs

@@ -0,0 +1,44 @@
+using JetBrains.Annotations;
+
+namespace CodeJam.Collections
+{
+	/// <summary>Suffix of the SuffixTree</summary>
+	[PublicAPI]
+	public struct Suffix
+	{
+		/// <summary>Buffer of all added strings</summary>
+		private readonly string _buffer;
+		/// <summary>Offset to the beginning of the suffix in the buffer</summary>
+		private readonly int _offset;
+
+		/// <summary>Constructs a new suffix</summary>
+		/// <param name="buffer">Buffer with all added strings</param>
+		/// <param name="sourceIndex">Source string index</param>
+		/// <param name="offset">Offset of the suffix in the buffer</param>
+		/// <param name="length">Length of the suffix</param>
+		internal Suffix([NotNull] string buffer, int sourceIndex, int offset, int length)
+		{
+			DebugCode.NotNull(buffer, nameof(buffer));
+			DebugCode.ValidIndex(sourceIndex, nameof(sourceIndex));
+			DebugCode.ValidIndexAndCount(offset, nameof(offset), length, nameof(length), buffer.Length);
+			_buffer = buffer;
+			SourceIndex = sourceIndex;
+			_offset = offset;
+			Length = length;
+		}
+
+		/// <summary>
+		/// The index of the source string in the order or addition to the Suffix tree
+		/// <remarks>0 - for the first added string, 1 - for the second, etc</remarks>
+		/// </summary>
+		public int SourceIndex { get; }
+		/// <summary>The length of the suffix</summary>
+		public int Length { get; }
+		/// <summary>The suffix value</summary>
+		public string Value => _buffer.Substring(_offset, Length);
+
+		/// <summary>String conversion operator</summary>
+		/// <param name="suffix">The suffix to convert</param>
+		public static implicit operator string(Suffix suffix) => suffix.Value;
+	}
+}

Experimental/src/Collections/SuffixTreeBase.cs

@@ -5,12 +5,12 @@
 
 namespace CodeJam.Collections
 {
-    public abstract class SuffixTreeBase
-    {
+	public abstract class SuffixTreeBase
+	{
 		/// <summary>Node alignment in Print output</summary>
 		private const int Align = 6;
 		/// <summary>Root node index</summary>
-	    protected const int RootNodeIndex = 0;
+		protected const int RootNodeIndex = 0;
 
 		/// <summary>Tree nodes</summary>
 		private readonly List<Node> _nodes;
@@ -21,27 +21,27 @@ public abstract class SuffixTreeBase
 		/// <summary>Adds a new node</summary>
 		/// <param name="node">A node to add</param>
 		/// <returns>Index of the node</returns>
-	    protected int AddNode(Node node)
-	    {
-		    var index = _nodes.Count;
+		protected int AddNode(Node node)
+		{
+			var index = _nodes.Count;
 			_nodes.Add(node);
-		    return index;
-	    }
+			return index;
+		}
 
 		/// <summary>Updates the node at the index</summary>
 		/// <param name="index">The index to update</param>
 		/// <param name="node">The new node value</param>
-	    protected void UpdateNode(int index, Node node) => _nodes[index] = node;
+		protected void UpdateNode(int index, Node node) => _nodes[index] = node;
 
-	    /// <summary>Gets a node at the index</summary>
+		/// <summary>Gets a node at the index</summary>
 		/// <param name="index">The index of the node</param>
 		/// <returns>The node</returns>
-	    protected Node GetNode(int index) => _nodes[index];
+		protected Node GetNode(int index) => _nodes[index];
 
 		/// <summary>Number of nodes</summary>
-	    protected int NodesCount => _nodes.Count;
+		protected int NodesCount => _nodes.Count;
 
-	    /// <summary>Concatenated input strings</summary>
+		/// <summary>Concatenated input strings</summary>
 		protected string InternalData { get; private set; }
 
 		/// <summary>List of end positions of added strings inside the InternalData</summary>
@@ -75,73 +75,141 @@ public void Add([NotNull]string data)
 			BuildFor(begin, InternalData.Length);
 		}
 
+		/// <summary>Enumerates all suffixes in the suffix tree</summary>
+		/// <remarks>May return suffixes with the same value of the they are present in different source strings</remarks>
+		/// <returns>The enumeration of all suffixes</returns>
+		[PublicAPI]
+		public IEnumerable<Suffix> AllSuffixes()
+		{
+			if (Root.IsLeaf) // Empty tree
+			{
+				yield break;
+			}
+
+			var branchStack = new Stack<BranchPoint>();
+			var branchPoint = new BranchPoint { Node = Root, EdgeIndex = 0 };
+			var length = 0;
+			for (;;)
+			{
+				var edge = GetNode(branchPoint.Node.Children[branchPoint.EdgeIndex]);
+				var edgeLength = edge.Length;
+				length += edgeLength;
+				if (!edge.IsTerminal)
+				{
+					branchPoint.Length = edgeLength;
+					branchStack.Push(branchPoint);
+					branchPoint = new BranchPoint { Node = edge, EdgeIndex = 0 };
+					continue;
+				}
+
+				// We have descended to a terminal edge. Let's produce a suffix
+				var end = edge.End;
+				var offset = end - length;
+				var sourceIndex = GetSourceIndexByEnd(end);
+				yield return new Suffix(InternalData, sourceIndex, offset, length);
+
+				// Move to the next suffix branch
+				for (;;)
+				{
+					length -= edgeLength;
+					var nextEdgeIndex = branchPoint.EdgeIndex + 1;
+					if (nextEdgeIndex < branchPoint.Node.Children.Count)
+					{
+						branchPoint.EdgeIndex = nextEdgeIndex;
+						break;
+					}
+					// There is no more branches on the current level
+					// Return to the previous level
+					if (branchStack.Count == 0)
+					{
+						// no more branches to visit
+						DebugCode.AssertState(length == 0, "Illegal final length. Check logic");
+						yield break;
+					}
+					branchPoint = branchStack.Pop();
+					edgeLength = branchPoint.Length;
+				}
+			}
+		}
+
+		/// <summary>Locates the source string index by the suffix end</summary>
+		/// <param name="end">The suffix end</param>
+		/// <returns>The source string index</returns>
+		private int GetSourceIndexByEnd(int end)
+		{
+			var index = EndPositions.LowerBound(end);
+			DebugCode.AssertState(index < EndPositions.Count && EndPositions[index] == end
+				, "Invalid source index computed. Check logic");
+			return index;
+		}
+
 		/// <summary>Appends suffixes for the last added string</summary>
-	    protected abstract void BuildFor(int begin, int end);
+		protected abstract void BuildFor(int begin, int end);
 
 		/// <summary>Creates a comparer for nodes against a char</summary>
 		/// <returns>The comparer</returns>
-	    protected Func<int, char, int> GetComparer() => (index, c) =>
-	    {
-		    var node = GetNode(index);
-		    if (node.Begin == node.End) // no char always less than any char
-		    {
-			    return -1;
-		    }
+		protected Func<int, char, int> GetComparer() => (index, c) =>
+		{
+			var node = GetNode(index);
+			if (node.Begin == node.End) // no char always less than any char
+			{
+				return -1;
+			}
 			var firstChar = InternalData[node.Begin];
-		    return firstChar - c;
-	    };
+			return firstChar - c;
+		};
 
 		/// <summary>Prints the tree structure to the string for the debugging purposes</summary>
 		/// <returns>The tree structure as a string</returns>
 		[Pure]
 		public string Print()
-	    {
-		    var sb = new StringBuilder();
-		    var currentIndex = RootNodeIndex;
-		    var stack = new List<ValueTuple<int, int>>();
-		    for (;;)
-		    {
-			    PrintNodeWithPath(sb, currentIndex, stack);
-			    var node = GetNode(currentIndex);
-			    if (node.Children != null)
-			    {
-				    stack.Add(ValueTuple.Create(currentIndex, node.Children.Count - 2));
-				    currentIndex = node.Children[node.Children.Count - 1];
+		{
+			var sb = new StringBuilder();
+			var currentIndex = RootNodeIndex;
+			var stack = new List<ValueTuple<int, int>>();
+			for (;;)
+			{
+				PrintNodeWithPath(sb, currentIndex, stack);
+				var node = GetNode(currentIndex);
+				if (node.Children != null)
+				{
+					stack.Add(ValueTuple.Create(currentIndex, node.Children.Count - 2));
+					currentIndex = node.Children[node.Children.Count - 1];
 					continue;
-			    }
-			    currentIndex = -1;
-			    while (stack.Count > 0)
-			    {
+				}
+				currentIndex = -1;
+				while (stack.Count > 0)
+				{
 					var t = stack[stack.Count - 1];
 					stack.RemoveAt(stack.Count - 1);
 					node = GetNode(t.Item1);
 					var nextChild = t.Item2;
-				    if (nextChild >= 0)
-				    {
-					    currentIndex = node.Children[nextChild];
+					if (nextChild >= 0)
+					{
+						currentIndex = node.Children[nextChild];
 						stack.Add(ValueTuple.Create(t.Item1, nextChild - 1));
-					    break;
-				    }
+						break;
+					}
 				}
-			    if (currentIndex == -1)
-			    {
-				    break;
-			    }
-		    }
+				if (currentIndex == -1)
+				{
+					break;
+				}
+			}
 			return sb.ToString();
-	    }
+		}
 
 		/// <summary>Prints a single node representation along with the path prefix</summary>
 		/// <param name="sb">The builder to print to</param>
 		/// <param name="nodeIndex">THe index of the node</param>
 		/// <param name="stack">The stack of nodes to process</param>
-	    private void PrintNodeWithPath([NotNull] StringBuilder sb, int nodeIndex
+		private void PrintNodeWithPath([NotNull] StringBuilder sb, int nodeIndex
 			, [NotNull] IReadOnlyList<ValueTuple<int, int>> stack)
-	    {
-		    if (stack.Count > 0)
-		    {
-			    for (var i = 0; i < stack.Count - 1; ++i)
-			    {
+		{
+			if (stack.Count > 0)
+			{
+				for (var i = 0; i < stack.Count - 1; ++i)
+				{
 					sb.Append(stack[i].Item2 >= 0 ? '|' : ' ');
 					sb.Append(' ', Align - 1);
 				}
@@ -155,13 +223,13 @@ private void PrintNodeWithPath([NotNull] StringBuilder sb, int nodeIndex
 				sb.Append('_', Align - 1);
 			}
 			PrintNodeText(sb, nodeIndex);
-	    }
+		}
 
 		/// <summary>Prints a single node information</summary>
 		/// <param name="sb">The builder to print to</param>
 		/// <param name="nodeIndex">The node index</param>
-	    protected virtual void PrintNodeText([NotNull] StringBuilder sb, int nodeIndex)
-	    {
+		protected virtual void PrintNodeText([NotNull] StringBuilder sb, int nodeIndex)
+		{
 			var n = GetNode(nodeIndex);
 			sb.AppendLine($"({nodeIndex}, [{n.Begin}-{n.End}), {InternalData.Substring(n.Begin, n.Length)})");
 		}
@@ -175,7 +243,7 @@ protected struct Node
 			/// <param name="begin">An edge start offset</param>
 			/// <param name="end">An edge end offset</param>
 			/// <param name="terminal">Is the edge terminates the string or not</param>
-			public Node(int begin, int end, bool terminal) : this(begin, end, terminal, null) {}
+			public Node(int begin, int end, bool terminal) : this(begin, end, terminal, null) { }
 
 			/// <summary>Constructs a new node</summary>
 			/// <param name="begin">An edge start offset</param>
@@ -206,5 +274,16 @@ public Node(int begin, int end, bool terminal, List<int> children)
 			/// <summary>Length of the corresponding substring</summary>
 			public int Length => End - Begin;
 		}
+
+		/// <summary>Branching point</summary>
+		private class BranchPoint
+		{
+			/// <summary>The tree node</summary>
+			public Node Node;
+			/// <summary>The chosen edge</summary>
+			public int EdgeIndex;
+			/// <summary>The length over the edge</summary>
+			public int Length;
+		}
 	}
 }

Experimental/tests/Collections/SuffixTreeTest.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Linq;
 
 using NUnit.Framework;
@@ -57,6 +58,42 @@ public void Test15RandomMultiple()
 			}
 		}
 
+		[Test]
+		public void Test16AllSuffixes()
+		{
+			const int length = 50;
+			for (var numberOfString = 1; numberOfString < 6; ++numberOfString)
+			{
+				var strings = Enumerable.Range(0, numberOfString)
+					.Select(_ => MakeRandomString(length)).ToArray();
+				var expectedSuffixes = new List<string>();
+				var expectedCounts = new LazyDictionary<string, List<int>>(_ => new List<int>());
+				var st = new SuffixTree();
+				for (var i = 0; i < strings.Length; ++i)
+				{
+					var s = strings[i];
+					st.Add(s);
+					for (var j = 0; j < s.Length; ++j)
+					{
+						var suffix = s.Substring(j);
+						expectedSuffixes.Add(suffix);
+						expectedCounts[suffix].Add(i);
+					}
+				}
+				st.Compact();
+				expectedSuffixes.Sort();
+				var suffixes = st.AllSuffixes().ToList();
+				Assert.That(suffixes.Select(_ => _.Value).ToList(), Is.EqualTo(expectedSuffixes));
+				var grouped = suffixes.Select(_ => new { value = _.Value, source = _.SourceIndex })
+					.GroupBy(_ => _.value).ToDictionary(_ => _.Key, _ => _.Select(v => v.source).OrderBy(v => v).ToList());
+				foreach (var v in grouped)
+				{
+					Assert.That(v.Value, Is.EqualTo(expectedCounts[v.Key]));
+				}
+			}
+		}
+
+
 		protected override void Check(string expected, params string[] data)
 		{
 			base.Check(expected, data);