Word Break Problem | (Trie solution)

Difficulty Level : Medium
Last Updated : 01 Aug, 2022

Read

Discuss

Practice

Video

Courses

Given an input string and a dictionary of words, find out if the input string can be segmented into a space-separated sequence of dictionary words. See following examples for more details.

This is a famous Google interview question, also being asked by many other companies now a days.

Consider the following dictionary 
{ i, like, sam, sung, samsung, mobile, ice, 
  cream, icecream, man, go, mango}

Input:  ilike
Output: Yes 
The string can be segmented as "i like".

Input:  ilikesamsung
Output: Yes
The string can be segmented as "i like samsung" or 
"i like sam sung".

The solution discussed here is mainly an extension of below DP based solution.
Dynamic Programming | Set 32 (Word Break Problem)

In the above post, a simple array is used to store and search words in a dictionary. Here we use Trie to do these tasks quickly.

Implementation:

C++

// A DP and Trie based program to test whether
// a given string can be segmented into
// space separated words in dictionary
#include <iostream>
using namespace std;
 
const int ALPHABET_SIZE = 26;
 
// trie node
struct TrieNode {
    struct TrieNode* children[ALPHABET_SIZE];
 
    // isEndOfWord is true if the node represents
    // end of a word
    bool isEndOfWord;
};
 
// Returns new trie node (initialized to NULLs)
struct TrieNode* getNode(void)
{
    struct TrieNode* pNode = new TrieNode;
 
    pNode->isEndOfWord = false;
 
    for (int i = 0; i < ALPHABET_SIZE; i++)
        pNode->children[i] = NULL;
 
    return pNode;
}
 
// If not present, inserts key into trie
// If the key is prefix of trie node, just
// marks leaf node
void insert(struct TrieNode* root, string key)
{
    struct TrieNode* pCrawl = root;
 
    for (int i = 0; i < key.length(); i++) {
        int index = key[i] - 'a';
        if (!pCrawl->children[index])
            pCrawl->children[index] = getNode();
 
        pCrawl = pCrawl->children[index];
    }
 
    // mark last node as leaf
    pCrawl->isEndOfWord = true;
}
 
// Returns true if key presents in trie, else
// false
bool search(struct TrieNode* root, string key)
{
    struct TrieNode* pCrawl = root;
 
    for (int i = 0; i < key.length(); i++) {
        int index = key[i] - 'a';
        if (!pCrawl->children[index])
            return false;
 
        pCrawl = pCrawl->children[index];
    }
 
    return (pCrawl != NULL && pCrawl->isEndOfWord);
}
 
// returns true if string can be segmented into
// space separated words, otherwise returns false
bool wordBreak(string str, TrieNode* root)
{
    int size = str.size();
 
    // Base case
    if (size == 0)
        return true;
 
    // Try all prefixes of lengths from 1 to size
    for (int i = 1; i <= size; i++) {
        // The parameter for search is str.substr(0, i)
        // str.substr(0, i) which is prefix (of input
        // string) of length 'i'. We first check whether
        // current prefix is in dictionary. Then we
        // recursively check for remaining string
        // str.substr(i, size-i) which is suffix of
        // length size-i
        if (search(root, str.substr(0, i))
            && wordBreak(str.substr(i, size - i), root))
            return true;
    }
 
    // If we have tried all prefixes and none
    // of them worked
    return false;
}
 
// Driver program to test above functions
int main()
{
    string dictionary[]
        = { "mobile", "samsung",  "sam",  "sung", "ma\n",
            "mango",  "icecream", "and",  "go",   "i",
            "like",   "ice",      "cream" };
    int n = sizeof(dictionary) / sizeof(dictionary[0]);
    struct TrieNode* root = getNode();
 
    // Construct trie
    for (int i = 0; i < n; i++)
        insert(root, dictionary[i]);
 
    wordBreak("ilikesamsung", root) ? cout << "Yes\n"
                                    : cout << "No\n";
    wordBreak("iiiiiiii", root) ? cout << "Yes\n"
                                : cout << "No\n";
    wordBreak("", root) ? cout << "Yes\n" : cout << "No\n";
    wordBreak("ilikelikeimangoiii", root) ? cout << "Yes\n"
                                          : cout << "No\n";
    wordBreak("samsungandmango", root) ? cout << "Yes\n"
                                       : cout << "No\n";
    wordBreak("samsungandmangok", root) ? cout << "Yes\n"
                                        : cout << "No\n";
    return 0;
}

Java

// A DP and Trie based program to test whether
// a given string can be segmented into
// space separated words in dictionary
import java.io.*;
import java.util.*;
 
class GFG {
 
    static final int ALPHABET_SIZE = 26;
 
    // trie node
    static class TrieNode {
        TrieNode children[];
 
        // isEndOfWord is true if the node
        // represents end of a word
        boolean isEndOfWord;
 
        // Constructor of TrieNode
        TrieNode()
        {
            children = new TrieNode[ALPHABET_SIZE];
            for (int i = 0; i < ALPHABET_SIZE; i++)
                children[i] = null;
 
            isEndOfWord = false;
        }
    }
 
    // If not present, inserts key into trie
    // If the key is prefix of trie node, just
    // marks leaf node
    static void insert(TrieNode root, String key)
    {
        TrieNode pCrawl = root;
 
        for (int i = 0; i < key.length(); i++) {
            int index = key.charAt(i) - 'a';
            if (pCrawl.children[index] == null)
                pCrawl.children[index] = new TrieNode();
 
            pCrawl = pCrawl.children[index];
        }
 
        // Mark last node as leaf
        pCrawl.isEndOfWord = true;
    }
 
    // Returns true if key presents in trie, else
    // false
    static boolean search(TrieNode root, String key)
    {
        TrieNode pCrawl = root;
 
        for (int i = 0; i < key.length(); i++) {
            int index = key.charAt(i) - 'a';
            if (pCrawl.children[index] == null)
                return false;
 
            pCrawl = pCrawl.children[index];
        }
        return (pCrawl != null && pCrawl.isEndOfWord);
    }
 
    // Returns true if string can be segmented
    // into space separated words, otherwise
    // returns false
    static boolean wordBreak(String str, TrieNode root)
    {
        int size = str.length();
 
        // Base case
        if (size == 0)
            return true;
 
        // Try all prefixes of lengths from 1 to size
        for (int i = 1; i <= size; i++) {
 
            // The parameter for search is
            // str.substring(0, i)
            // str.substring(0, i) which is
            // prefix (of input string) of
            // length 'i'. We first check whether
            // current prefix is in dictionary.
            // Then we recursively check for remaining
            // string str.substr(i, size) which
            // is suffix of length size-i.
            if (search(root, str.substring(0, i))
                && wordBreak(str.substring(i, size), root))
                return true;
        }
 
        // If we have tried all prefixes and none
        // of them worked
        return false;
    }
 
    // Driver code
    public static void main(String[] args)
    {
        String dictionary[]
            = { "mobile", "samsung",  "sam",  "sung", "ma",
                "mango",  "icecream", "and",  "go",   "i",
                "like",   "ice",      "cream" };
 
        int n = dictionary.length;
        TrieNode root = new TrieNode();
 
        // Construct trie
        for (int i = 0; i < n; i++)
            insert(root, dictionary[i]);
 
        System.out.print(wordBreak("ilikesamsung", root)
                             ? "Yes\n"
                             : "No\n");
        System.out.print(
            wordBreak("iiiiiiii", root) ? "Yes\n" : "No\n");
        System.out.print(wordBreak("", root) ? "Yes\n"
                                             : "No\n");
        System.out.print(
            wordBreak("ilikelikeimangoiii", root) ? "Yes\n"
                                                  : "No\n");
        System.out.print(wordBreak("samsungandmango", root)
                             ? "Yes\n"
                             : "No\n");
        System.out.print(wordBreak("samsungandmangok", root)
                             ? "Yes\n"
                             : "No\n");
    }
}
 
// This code is contributed by Ganeshchowdharysadanala

Python3

class Solution(object):
    def wordBreak(self, s, wordDict):
        """
        Author : @amitrajitbose
        :type s: str
        :type wordDict: List[str]
        :rtype: bool
        """
        """CREATING THE TRIE CLASS"""
 
        class TrieNode(object):
 
            def __init__(self):
                self.children = []  # will be of size = 26
                self.isLeaf = False
 
            def getNode(self):
                p = TrieNode()  # new trie node
                p.children = []
                for i in range(26):
                    p.children.append(None)
                p.isLeaf = False
                return p
 
            def insert(self, root, key):
                key = str(key)
                pCrawl = root
                for i in key:
                    index = ord(i)-97
                    if(pCrawl.children[index] == None):
                        # node has to be initialised
                        pCrawl.children[index] = self.getNode()
                    pCrawl = pCrawl.children[index]
                pCrawl.isLeaf = True  # marking end of word
 
            def search(self, root, key):
                # print("Searching %s" %key) #DEBUG
                pCrawl = root
                for i in key:
                    index = ord(i)-97
                    if(pCrawl.children[index] == None):
                        return False
                    pCrawl = pCrawl.children[index]
                if(pCrawl and pCrawl.isLeaf):
                    return True
 
        def checkWordBreak(strr, root):
            n = len(strr)
            if(n == 0):
                return True
            for i in range(1, n+1):
                if(root.search(root, strr[:i]) and checkWordBreak(strr[i:], root)):
                    return True
            return False
 
        """IMPLEMENT SOLUTION"""
        root = TrieNode().getNode()
        for w in wordDict:
            root.insert(root, w)
        out = checkWordBreak(s, root)
        if(out):
            return "Yes"
        else:
            return "No"
 
 
print(Solution().wordBreak("thequickbrownfox",
                           ["the", "quick", "fox", "brown"]))
print(Solution().wordBreak("bedbathandbeyond", [
      "bed", "bath", "bedbath", "and", "beyond"]))
print(Solution().wordBreak("bedbathandbeyond", [
      "teddy", "bath", "bedbath", "and", "beyond"]))
print(Solution().wordBreak("bedbathandbeyond", [
      "bed", "bath", "bedbath", "and", "away"]))

C#

// C# program to implement
// a DP and Trie based program to test whether
// a given string can be segmented into
// space separated words in dictionary
using System;
 
class GFG {
 
  static readonly int ALPHABET_SIZE = 26;
 
  // trie node
  class TrieNode {
    public TrieNode[] children;
 
    // isEndOfWord is true if the node
    // represents end of a word
    public bool isEndOfWord;
 
    // Constructor of TrieNode
    public TrieNode()
    {
      children = new TrieNode[ALPHABET_SIZE];
      for (int i = 0; i < ALPHABET_SIZE; i++)
        children[i] = null;
 
      isEndOfWord = false;
    }
  }
 
  // If not present, inserts key into trie
  // If the key is prefix of trie node, just
  // marks leaf node
  static void insert(TrieNode root, string key)
  {
    TrieNode pCrawl = root;
 
    for (int i = 0; i < key.Length; i++) {
      int index = key[i] - 'a';
      if (pCrawl.children[index] == null)
        pCrawl.children[index] = new TrieNode();
 
      pCrawl = pCrawl.children[index];
    }
 
    // Mark last node as leaf
    pCrawl.isEndOfWord = true;
  }
 
  // Returns true if key presents in trie, else
  // false
  static bool search(TrieNode root, string key)
  {
    TrieNode pCrawl = root;
 
    for (int i = 0; i < key.Length; i++) {
      int index = key[i] - 'a';
      if (pCrawl.children[index] == null)
        return false;
 
      pCrawl = pCrawl.children[index];
    }
    return (pCrawl != null && pCrawl.isEndOfWord);
  }
 
  // Returns true if string can be segmented
  // into space separated words, otherwise
  // returns false
  static bool wordBreak(string str, TrieNode root)
  {
    int size = str.Length;
 
    // Base case
    if (size == 0)
      return true;
 
    // Try all prefixes of lengths from 1 to size
    for (int i = 1; i <= size; i++) {
      // The parameter for search is
      // str.Substring(0, i) which is
      // prefix (of input string) of
      // length 'i'. We first check whether
      // current prefix is in dictionary.
      // Then we recursively check for remaining
      // string str.Substring(i, size-i) which
      // is suffix of length size-i.
      if (search(root, str.Substring(0, i))
          && wordBreak(str.Substring(i, size - i),
                       root))
        return true;
    }
 
    // If we have tried all prefixes and none
    // of them worked
    return false;
  }
 
  // Driver code
  static void Main(string[] args)
  {
    string[] dictionary = new string[] {
      "mobile", "samsung",  "sam",  "sung", "ma",
      "mango",  "icecream", "and",  "go",   "i",
      "like",   "ice",      "cream"
      };
 
    int n = dictionary.Length;
    TrieNode root = new TrieNode();
 
    // Construct trie
    for (int i = 0; i < n; i++)
      insert(root, dictionary[i]);
 
    Console.WriteLine(
      wordBreak("ilikesamsung", root) ? "Yes" : "No");
    Console.WriteLine(
      wordBreak("iiiiiiii", root) ? "Yes" : "No");
    Console.WriteLine(wordBreak("", root) ? "Yes"
                      : "No");
    Console.WriteLine(
      wordBreak("ilikelikeimangoiii", root) ? "Yes"
      : "No");
    Console.WriteLine(wordBreak("samsungandmango", root)
                      ? "Yes"
                      : "No");
    Console.WriteLine(
      wordBreak("samsungandmangok", root) ? "Yes"
      : "No");
  }
}
 
// This code is contributed by cavi4762.

Javascript

// A DP and Trie based program to test whether
// a given string can be segmented into
// space separated words in dictionary
 
let ALPHABET_SIZE = 26;
 
// trie node
class TrieNode {
    constructor()
    {
        this.children = new Array(ALPHABET_SIZE);
 
        // isEndOfWord is true if the node represents
        // end of a word
        this.isEndOfWord;
    }
};
 
// Returns new trie node (initialized to NULLs)
function getNode()
{
    let pNode = new TrieNode;
 
    pNode.isEndOfWord = false;
 
    for (var i = 0; i < ALPHABET_SIZE; i++)
        pNode.children[i] = null;
 
    return pNode;
}
 
// If not present, inserts key into trie
// If the key is prefix of trie node, just
// marks leaf node
function insert(root, key)
{
    let pCrawl = root;
 
    for (let i = 0; i < key.length; i++) {
        let index
            = key[i].charCodeAt(0) - 'a'.charCodeAt(0);
        if (pCrawl.children[index] == null)
            pCrawl.children[index] = getNode();
 
        pCrawl = pCrawl.children[index];
    }
 
    // mark last node as leaf
    pCrawl.isEndOfWord = true;
}
 
// Returns true if key presents in trie, else
// false
function search(root, key)
{
    let pCrawl = root;
 
    for (let i = 0; i < key.length; i++) {
        let index
            = key[i].charCodeAt(0) - 'a'.charCodeAt(0);
        if (pCrawl.children[index] == null)
            return false;
 
        pCrawl = pCrawl.children[index];
    }
 
    return (pCrawl != null && pCrawl.isEndOfWord);
}
 
// returns true if string can be segmented into
// space separated words, otherwise returns false
function wordBreak(str, root)
{
    let size = str.length;
 
    // Base case
    if (size == 0)
        return true;
 
    // Try all prefixes of lengths from 1 to size
    for (let i = 1; i <= size; i++) {
        // The parameter for search is str.substr(0, i)
        // str.substr(0, i) which is prefix (of input
        // string) of length 'i'. We first check whether
        // current prefix is in dictionary. Then we
        // recursively check for remaining string
        // str.substr(i, size-i) which is suffix of
        // length size-i
        if (search(root, str.substr(0, i))
            && wordBreak(str.substr(i, size - i), root))
            return true;
    }
 
    // If we have tried all prefixes and none
    // of them worked
    return false;
}
 
// Driver program to test above functions
let dictionary = [
    "mobile", "samsung", "sam", "sung", "ma\n", "mango",
    "icecream", "and", "go", "i", "like", "ice", "cream"
];
let n = dictionary.length;
let root = getNode();
 
// Construct trie
for (let i = 0; i < n; i++)
    insert(root, dictionary[i]);
 
wordBreak("ilikesamsung", root) ? console.log("Yes")
                                : console.log("No");
wordBreak("iiiiiiii", root) ? console.log("Yes")
                            : console.log("No");
wordBreak("ilikelikeimangoiii", root) ? console.log("Yes")
                                      : console.log("No");
wordBreak("samsungandmango", root) ? console.log("Yes")
                                   : console.log("No");
wordBreak("samsungandmangok", root) ? console.log("Yes")
                                    : console.log("No");
 
// This code is contributed by phasing17

Output

Yes
Yes
Yes
Yes
Yes
No

This article is contributed by Pranav. The Python code has been contributed by Jeet9. If you like GeeksforGeeks and would like to contribute, you can also write an article using write.geeksforgeeks.org or mail your article to review-team@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

My Personal Notes