java 677. Map Sum Pairs（#Trie）.java

Posted 2021-05-08

/** Initialize your data structure here. */
Map<String, Integer> map;
Map<String, Integer> original;
public MapSum() {
    map = new HashMap<>();
    original = new HashMap<>();
}

public void insert(String key, int val) {
    val -= original.getOrDefault(key, 0); // calculate the diff to be added to prefixes
    String s = "";
    for(char c : key.toCharArray()) {
        s += c; // creating all prefixes
        map.put(s, map.getOrDefault(s, 0) + val); //update/insert all prefixes with new value
    }
    original.put(key, original.getOrDefault(key, 0) + val);
}

public int sum(String prefix) {
    return map.getOrDefault(prefix, 0);
}

class TrieNode{
    
    private TrieNode[] children;
    public boolean hasWord;
    public int sum;
    public Map<String, Integer> map;
    
    public TrieNode() {
        children = new TrieNode[26];
        hasWord = false;
    }
    
    public void insert(String word, int index, int val){
        if( index == word.length()){
            hasWord = true;
            return; 
        }
        
        int pos = word.charAt(index) - 'a';
        
        if(children[pos] == null){
            children[pos] = new TrieNode();
        }  
        
        if (children[pos].map.containsKey(word)) {
            children[pos].sum += val - children[pos].map.get(word);
            
        } else {
            children[pos].sum += val;
            
        }
        children[pos].insert(word, index + 1, val);
    }
    
    public TrieNode find(String word, int index){
        if( index == word.length()){
            return this;
        }
        
        int pos = word.charAt(index) - 'a';
        
        if(children[pos] == null){
            return null;
        } 
        
        return children[pos].find(word, index + 1);
        
    }
}

public class MapSum {

    /** Initialize your data structure here. */
    TrieNode root;
    public MapSum() {
        root = new TrieNode();
    }
    
    public void insert(String key, int val) {
        root.insert(key, 0, val);
    }
    
    public int sum(String prefix) {
        TrieNode node = root.find(prefix, 0);
        return node == null ? 0 : node.sum;
    }
}

/**
 * Your MapSum object will be instantiated and called as such:
 * MapSum obj = new MapSum();
 * obj.insert(key,val);
 * int param_2 = obj.sum(prefix);
 */