String Tokenization Problem

Recently, I ran into a LeetCode-style problem from an Anthropic interview, and I wanted to share my solution along with my step-by-step optimization process.

Problem definition can be found **here.**

Requirements and the Brute-force Solution

Note: the website only offers solutions in either Python, Java, TypeScript or C++, but for my own convenience and preference I have written it in Go.

As stated in the problem description, the tokenization process must follow these rules:

• Longest Match Priority: The longest key from the dictionary must be chosen for the output.
• Greedy Consumption: Continue processing from the last consumed position.
• Literal Preservation: If no dictionary key matches at the current position, the unmatched character should be preserved as an individual literal token.
• Output Format: Output IDs that are corresponding to the found keys as a strings array.

Constraints:

• 1 ≤ text.length ≤ 10⁹
• 0 ≤ dictionary.length ≤ 10⁹
• All tokens in dictionary are unique.

Having everything laid out, let’s try to find them most naive solution.

Firstly, let’s define the function and the variables:

func tokenize(text string, dictionary map[string]int) {
    output := []string{}
    pos := 0 

    // algorithm

    return output    
}

output: This is our slice of result strings. Nothing fancy here.

pos: This pointer tracks our current index in the input string, allowing us to resume matching immediately after our last consumed token.

While writing these variables down, the first algorithm that popped up in my head was to loop through dictionary keys and compare them to slices of the input string using pos variable. This approach looked like this:

func tokenize(text string, dictionary map[string]int) {
    output := []string{}
    pos := 0 

    for k, i := range dictionary {
        for _, c := range text {
            slice := text[pos:pos+len(c)] 

            if slice == k {
                output = append(output, strconv.Itoa(i)) 
                pos += len(c) 
            }
        }  
    }

    return output    
}

I quickly realized that it is not what I was looking for as my final solution. But the initial logic indeed ended up being something similar, just with some corrections.

For the Longest Matching Priority I needed to sort the dictionary keys in descending order by their length.

skeys := slices.Collect(maps.Keys(dictionary))
lenCmp := func(a, b string) int {
    return cmp.Compare(len(b), len(a))
}

slices.SortFunc(skeys, lenCmp)

Ok, now we have something to work with. The slice calculated can be out of bounds, so let’s fix it:

for pos < len(text) {
    for _, k := range skeys {
        if pos+len(k) <= len(text) {
            slice := text[pos:pos+len(k)] 

            if slice == k {
                id := strconv.Itoa(dictionary[k])
                output = append(output, id)
                pos += len(k) 
                break
            } 
        }
    } 
}

Now it is almost a working solution, but how do we spell the remaining not matched tokens by characters? The answer is obvious, just track whether a second if statement worked with the matched variable.

 for pos < len(text) {
    matched := false

    for _, k := range skeys {
        if pos+len(k) <= len(text) {
            slice := text[pos : pos+len(k)]

            if slice == k {
                id := strconv.Itoa(dictionary[k])
                output = append(output, id)
                pos += len(k)
                matched = true
                break
            }
        }
    }

    if !matched {
        output = append(output, string(text[pos]))
         pos += 1
    }
}

The full working solution:

func tokenize(text string, dictionary map[string]int) []string {
    output := []string{}
    pos := 0

    skeys := slices.Collect(maps.Keys(dictionary))
    lenCmp := func(a, b string) int {
        return cmp.Compare(len(b), len(a))
    }

    slices.SortFunc(skeys, lenCmp)

    for pos < len(text) {
        matched := false

        for _, k := range skeys {
            if pos+len(k) <= len(text) {
                slice := text[pos : pos+len(k)]

                if slice == k {
                    id := strconv.Itoa(dictionary[k])
                    output = append(output, id)
                    pos += len(k)
                    matched = true
                    break
                }
            }
        }

        if !matched {
            output = append(output, string(text[pos]))
             pos += 1
        }
    }

    return output
}

Time and space complexity: O(N M L) and O(N + M L) *correspondingly.

Trie Solution(Prefix tree)

There was clearly a room for a structural optimization, and a Prefix Tree(Trie) was the perfect fit.

Trie(also known as Prefix tree) is a specialized tree-based data structure used to efficiently store and retrieve keys in a dataset of strings.

The visualization looks like this:

This alone can drastically improve out application performance and scalability. This is how Trie implementation looks like in Go:

package main

import "fmt"

// Node structure and a flag for the end of the word.
type Node struct {
    children map[rune]*Node // Child nodes
    isEnd    bool           // Indicates if it's the end of a word
    id       int            // We will need it for our solution
}

// NewNode creates and initializes a new node
func NewNode(id int) *Node {
    return &Node{children: make(map[rune]*Node), id: id}
}

// Trie represents the Trie structure
// containing a pointer to the root node
type Trie struct {
    root *Node 
}

// InitTrie initializes the trie structure and return its pointer.
func InitTrie() *Trie {
    return &Trie{root: NewNode(0)}
}

// Insert adds a word to the Trie
func (t *Trie) Insert(word string, id int) {
    node := t.root // Start at the root
    for _, char := range word {
        if _, exists := node.children[char]; !exists {
            node.children[char] = NewNode(id) // Create a new node
        }
        node = node.children[char] // Move to the child node
    }
    node.isEnd = true // Mark the end of the word
}

// Search checks if a word exists in the Trie
func (t *Trie) Search(word string) bool {
    node := t.root // Start at the root
    for _, char := range word { 
        if _, exists := node.children[char]; !exists { 
            return false // Word not found
        }
        node = node.children[char] // Move to the child node
    }
    return node.isEnd // Return true if it's an end of a word
}

For out purposes, we will need to slightly modify a search function. The algorithm stays similar to the initial solution, but we are looking for an end node, which will signal us that the word has been found. Before that, the final script initialized a tree with nodes in it, which correspond to dictionary keys and IDs. Take a look at it:

func tokenize(text string, dictionary map[string]int) []string {
   output := []string{}
   trie := InitTrie()

   for k, i := range dictionary {
      trie.Insert(k, i)
   }

   pos := 0
   for pos < len(text) {
      node := trie.root
      longestMatchLen := 0
      matchedId := ""

      for i := pos; i < len(text); i++ {
          c := rune(text[i])

          nextNode, exists := node.children[c]

          if !exists {
              break
          }
          node = nextNode

          if node.isEnd {
              longestMatchLen = (i - pos) + 1
              matchedId = strconv.Itoa(node.id)
          }
      }

      if longestMatchLen > 0 {
          output = append(output, matchedId)
          pos += longestMatchLen
      } else {
          output = append(output, string(text[pos]))
          pos++
      }
   }

   return output
}

Now we have a much more scalable solution with O(N L) time complexity and O(M L ∑) *space complexity.

Summary

As stated in the website that I mentioned in this article, this problem appears in Anthropic’s and Google’s interview rounds. While gathering information about Trie data structure, I also came up upon different use cases of it, for example a search engine.

The two solutions presented here can be found in this Github repo under string-tokenization directory.