LeetCode/java/replace-words.java

import java.util.List;

class Solution {
  private class Trie {
    public Trie() {
      root = new TrieNode();
    }

    public void insert(String word) {
      root.insert(word);
    }

    public String shortestPrefix(String word) {
      return root.shortestPrefix(word);
    }

    private TrieNode root;

    private class TrieNode {
      boolean has;
      TrieNode[] successors;

      TrieNode() {
        has = false;
        successors = new TrieNode[26];
      }

      void insert(String word) {
        TrieNode node = this;

        for (int i = 0; i < word.length(); ++i) {
          int idx = word.charAt(i) - 'a';

          if (node.successors[idx] == null) {
            node.successors[idx] = new TrieNode();
          }
          node = node.successors[idx];
        }

        node.has = true;
      }

      String shortestPrefix(String word) {
        TrieNode node = this;

        for (int i = 0; i < word.length(); ++i) {
          int idx = word.charAt(i) - 'a';

          if (node.successors[idx] == null) {
            return word;
          }

          node = node.successors[idx];
          if (node.has) {
            return word.substring(0, i + 1);
          }
        }

        return word;
      }
    }
  }

  public String replaceWords(List<String> dictionary, String sentence) {
    String words[] = sentence.split(" ");

    var t = new Trie();
    for (var word : dictionary) {
      t.insert(word);
    }

    for (int i = 0; i < words.length; ++i) {
      words[i] = t.shortestPrefix(words[i]);
    }

    return String.join(" ", words);
  }
}