Trie Data Structure Advanced Concepts 1

Photo by Sean Pollock on Unsplash

A rare, yet useful data structure, helpful especially in problems associated to prefix and suffix.

If you know the basic functions of Tries, you can understand how to solve the advanced question related to prefix and suffix (given below) using Tries.

Let us consider the question:

Given a list of N words denoted by string array A of size N.

You must answer Q queries denoted by string array B, for each query you have a string S of size M, find the number of words in the list A which have string S as a prefix and suffix.

NOTE: The size M for all strings in the queries remains the same.

Return an integer array of size Q denoting the answer of each query.

A = [“ababa”, “aabbvaab”, “aecdsaaec”, “abaaxbqaba”]

B = [“aba”, “aec”, “abb”, “aab”]

Output: [2, 1, 0, 1]

Explanation:

2 word “ababa” and “abaaxbqaba” has both prefix and suffix “aba”.

1 word “aecdsaaec” has both prefix and suffix “aec”.

No word has both prefix and suffix “abb”.

1 word “aabbvaab” has both prefix and suffix “aab”.

Most people start sweating in interviews when given such a question.

My mentor, currently working as an SDE in Google, asked me to begin from the logic and then build, and code the required functions. This is the best approach to solve competitive questions quickly. Let’s begin by initializing the Trie Data Structure.

Now let us begin with the logic:

Points of Focus:

1. String S as a prefix and suffix.
2. Size M for all strings in the queries remains the same.

Common Sense derivations:

1. In any case, for an answer to be valid, the first and the last M letters of a word in array A should be the same i.e. A[0: M] == [A.length() — M: A.length()]
2. If we store the words in a Trie Data Structure that qualify the above condition, we won’t need to check for the prefix as well as the suffix. Just checking the prefix shall do the work since the prefix will be the same as the suffix.

Basic Trie Data Structure

Basic logic to insert in Trie

Now that we have coded the logic of sorting out the words that are to be inserted in the Trie, we need to define the insert function i.e., how shall we insert the strings in the Trie.

Basic Function to insert strings in Trie

Till this point: We have a Trie Data Structure that contains words that have the same prefix and suffix of size M.

Now all that’s left is answering all the queries.

The Logic:

We shall count the number of words having the prefix == query. Remember Trie only contains words with prefix == suffix, so we don’t need to worry about the suffix part of the problem.

Till now we haven’t stored the count of the words anywhere in the Trie. So let us make changes to the previously written code to add a count for each node in the Trie.

Trie structure with count

Let’s implement the increment of count while inserting the words in the Trie.

Insert in Trie with count increment

Finally, our main function would look something like this:

The logic of the solution

To return the count of strings with the queried prefix (count_prefix_strings), we just access the count of the node corresponding to the last letter of the queried prefix.

Returning the count parameter of the required node

If you were able to make it to the end, Congrats. You just understood one of the most basic, yet tricky applications of the Trie Data Structures.