triemap: Add GetLongestKeyValuePairs(), GetShortestKeyValuePairs()

rmandvikar · Nov 3, 2023 · 4104c08 · 4104c08
1 parent 0d31593
commit 4104c08
Show file tree

Hide file tree

Showing 4 changed files with 132 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -166,6 +166,12 @@ void Remove(string key);
 // Removes key prefix from TrieMap and return true else false.
 bool RemoveKeyPrefix(string keyPrefix);
 
+// Gets string->TValue pairs for longest keys from the Trie.
+IEnumerable<KeyValuePair<string, TValue>> GetLongestKeyValuePairs();
+
+// Gets string->TValue pairs for shortest keys from the Trie.
+IEnumerable<KeyValuePair<string, TValue>> GetShortestKeyValuePairs();
+
 // Gets string->TValue pair for longest prefix matching the word from the Trie.
 KeyValuePair<string, TValue>? GetLongestPrefixMatch(string word);
 
@@ -220,6 +226,12 @@ trieMap.Remove("key"); // throws
 // Removes by key prefix.
 bool isKeyPrefixRemoved = trieMap.RemoveKeyPrefix("key"); // isKeyPrefixRemoved = true
 
+// Gets string->TValue pairs for longest keys from the Trie.
+IEnumerable<KeyValuePair<string, TValue>> longestKeyValuePairs = trieMap.GetLongestKeyValuePairs();
+
+// Gets string->TValue pairs for shortest keys from the Trie.
+IEnumerable<KeyValuePair<string, TValue>> shortestKeyValuePairs = trieMap.GetShortestKeyValuePairs();
+
 // Get key-value item for longest prefix matching the word.
 KeyValuePair<string, TValue>? kvPairForKey12 = trieMapGetLongestPrefixMatch("key12345");
 

diff --git a/src/rm.Trie/ITrieMap.TValue.cs b/src/rm.Trie/ITrieMap.TValue.cs
@@ -74,6 +74,16 @@ public interface ITrieMap<TValue>
 	/// </summary>
 	bool RemoveKeyPrefix(string keyPrefix);
 
+	/// <summary>
+	/// Gets string->TValue pairs for longest keys from the Trie.
+	/// </summary>
+	IEnumerable<KeyValuePair<string, TValue>> GetLongestKeyValuePairs();
+
+	/// <summary>
+	/// Gets string->TValue pairs for shortest keys from the Trie.
+	/// </summary>
+	IEnumerable<KeyValuePair<string, TValue>> GetShortestKeyValuePairs();
+
 	/// <summary>
 	/// Gets string->TValue pair for longest prefix matching the word from the Trie.
 	/// </summary>

diff --git a/src/rm.Trie/TrieMap.TValue.cs b/src/rm.Trie/TrieMap.TValue.cs
@@ -204,6 +204,30 @@ public bool RemoveKeyPrefix(string keyPrefix)
 		return trieNode != null;
 	}
 
+	/// <summary>
+	/// Gets string->TValue pairs for longest keys from the Trie.
+	/// </summary>
+	public IEnumerable<KeyValuePair<string, TValue>> GetLongestKeyValuePairs()
+	{
+		var longestKeyValuePairs = new List<KeyValuePair<string, TValue>>();
+		var buffer = new StringBuilder();
+		var length = new Wrapped<int>(0);
+		GetLongestKeyValuePairs(rootTrieNode, longestKeyValuePairs, buffer, length);
+		return longestKeyValuePairs;
+	}
+
+	/// <summary>
+	/// Gets string->TValue pairs for shortest keys from the Trie.
+	/// </summary>
+	public IEnumerable<KeyValuePair<string, TValue>> GetShortestKeyValuePairs()
+	{
+		var shortestKeyValuePairs = new List<KeyValuePair<string, TValue>>();
+		var buffer = new StringBuilder();
+		var length = new Wrapped<int>(int.MaxValue);
+		GetShortestKeyValuePairs(rootTrieNode, shortestKeyValuePairs, buffer, length);
+		return shortestKeyValuePairs;
+	}
+
 	/// <summary>
 	/// Gets string->TValue pair for longest prefix matching the word from the Trie.
 	/// </summary>
@@ -262,6 +286,52 @@ private IEnumerable<TResult> Traverse<TResult>(TrieNode<TValue> trieNode,
 		}
 	}
 
+	private void GetLongestKeyValuePairs(TrieNode<TValue> trieNode,
+		ICollection<KeyValuePair<string, TValue>> longestKeyValuePairs, StringBuilder buffer, Wrapped<int> length)
+	{
+		if (trieNode.HasValue())
+		{
+			if (buffer.Length > length.Value)
+			{
+				longestKeyValuePairs.Clear();
+				length.Value = buffer.Length;
+			}
+			if (buffer.Length == length.Value)
+			{
+				longestKeyValuePairs.Add(new KeyValuePair<string, TValue>(buffer.ToString(), trieNode.Value));
+			}
+		}
+		foreach (var child in trieNode.GetChildren())
+		{
+			buffer.Append(child.Character);
+			GetLongestKeyValuePairs(child, longestKeyValuePairs, buffer, length);
+			buffer.Length--;
+		}
+	}
+
+	private void GetShortestKeyValuePairs(TrieNode<TValue> trieNode,
+		ICollection<KeyValuePair<string, TValue>> shortestKeyValuePairs, StringBuilder buffer, Wrapped<int> length)
+	{
+		if (trieNode.HasValue())
+		{
+			if (buffer.Length < length.Value)
+			{
+				shortestKeyValuePairs.Clear();
+				length.Value = buffer.Length;
+			}
+			if (buffer.Length == length.Value)
+			{
+				shortestKeyValuePairs.Add(new KeyValuePair<string, TValue>(buffer.ToString(), trieNode.Value));
+			}
+		}
+		foreach (var child in trieNode.GetChildren())
+		{
+			buffer.Append(child.Character);
+			GetShortestKeyValuePairs(child, shortestKeyValuePairs, buffer, length);
+			buffer.Length--;
+		}
+	}
+
 	private KeyValuePair<string, TValue>? GetLongestPrefixMatch(TrieNode<TValue> trieNode,
 		string word)
 	{

diff --git a/tests/rm.TrieTest/TrieMapTest.cs b/tests/rm.TrieTest/TrieMapTest.cs
@@ -248,6 +248,46 @@ public void RemoveKeyPrefix02()
 		Assert.AreEqual(0, trieO.Values().Count());
 	}
 
+	[Test]
+	public void GetLongestKeyValuePairs01()
+	{
+		var trie = BuildSampleTrie();
+		trie.Add("the longest word", 1337);
+		var expected = new[] { new KeyValuePair<string, int>("the longest word", 1337) };
+		var longestKeyValuePairs = trie.GetLongestKeyValuePairs();
+		Assert.AreEqual(expected, longestKeyValuePairs);
+	}
+
+	[Test]
+	public void GetLongestKeyValuePairs02()
+	{
+		var trie = BuildSampleTrie();
+		trie.Add("the longest word1", 1337);
+		trie.Add("the longest word2", 1337);
+		var expected = new[] { new KeyValuePair<string, int>("the longest word1", 1337), new KeyValuePair<string, int>("the longest word2", 1337) };
+		var longestKeyValuePairs = trie.GetLongestKeyValuePairs();
+		Assert.AreEqual(expected, longestKeyValuePairs);
+	}
+
+	[Test]
+	public void GetShortestKeyValuePairs01()
+	{
+		var trie = BuildSampleTrie();
+		var expected = new[] { new KeyValuePair<string, int>("key1", 1) };
+		var shortestKeyValuePairs = trie.GetShortestKeyValuePairs();
+		Assert.AreEqual(expected, shortestKeyValuePairs);
+	}
+
+	[Test]
+	public void GetShortestKeyValuePairs02()
+	{
+		var trie = BuildSampleTrie();
+		trie.Add("keya", 1337);
+		var expected = new[] { new KeyValuePair<string, int>("key1", 1), new KeyValuePair<string, int>("keya", 1337) };
+		var shortestKeyValuePairs = trie.GetShortestKeyValuePairs();
+		Assert.AreEqual(expected, shortestKeyValuePairs);
+	}
+
 	[Test]
 	public void GetLongestPrefixMatch01()
 	{