Google Additional Round 1 - Word Predictor

/*
Want to build a word predictor.

Given training data like this:
[
  ["I", "am", "Sam"],
  ["Sam", "I", "am"],
  ["I", "like", "green", "eggs", "and", "ham"],
]

Goal: Given a word A, predict the next word B as the most frequent word occurring after A.
Example: "I" has "am" after it twice, and "like" after it once, so predNextWord("I") = "am". 

"I" -> ["am" : 2, "like" : 1]
"am" -> ["Sam" : 1]


Time Complexity:
N = number of words in training data

WordPrector intiator: O(N)
LookUp(): O(N)


Space Complexity:
N = number of words in training data

WordPrector intiator: O(N)
LookUp(): O(N)



  ["I", "am", "I", "a", "I", "b", "I", "c"]

*/

class WorldPredictor {
  
  Map<String, Map<String, Integer>> dictionary;
  
  Map<String, String> cache;
  
  public WorldPredictor(String[][] data) {
    this.dictionary = new HashMap<>();
    this.cache = new HashMap<>();
    
    for (int i = 0; i < data.length; i++) {
      for (int j = 0; j < data[i].length-1; j++) {
        String word = data[i][j];
        String nextWord = data[i][j+1];
        Map<String, Integer> wordCountMap = dictionary.getOrDefault(word, new HashMap<>());
        
        int wordCount = wordCountMap.getOrDefault(nextWord, 0);
        wordCount++;
        
        wordCountMap.put(nextWord, wordCount);
        dictionary.put(word, wordcountMap);
      }
    }
    
    // create cache here
  }
  
  public String lookUp(String A) {
    Map<String, Integer> wordCountMap = dictionary.getOrDefault(A, new HashMap<>());
    int max = 0;
    String mostFrequentWord = "";
    
    if (cache.containsKey(A)) {
      return cache.get(A);
    }
    
    for (Map.Entry<String, Integer> entry : wordCountMap.EntrySet()) {
      if (entry.value() > max) {
        max = entry.value();
        mostFrequentWord = entry.key();
      }
    }
    
    cache.put(A, mostFrequentWord);
    
    return mostFrequentWord;
  }
  
}