Java O(length(key))-insert/sum Trie solution with clear explanation

  • 2

    The solution is based on the standard Trie and its insert and search methods, with changes for the fields of TrieNode. TrieNode is defined as TrieNode {int val; int sum; boolean isAWord; TrieNode[] children;}.

    To insert a pair {key, val}, in addition to the standard insert method, we update the sum as sum+=val of all the nodes along the insertion path. Notice that if the key existed in the Trie, we call insert again with {key, -old_val} so as to update the sum of all the nodes along the path again.

    To get the sum, we just apply standard search method but return the sum.

    class MapSum {
    	private TrieNode root;
    	/** Initialize your data structure here. */
    	public MapSum() {
    		root = new TrieNode(0);
    	public void insert(String key, int val) {
    		TrieNode node = root;
    		int ci;
    		for (int i = 0; i < key.length(); i++) {
    			ci = key.charAt(i) - 'a';
    			if (node.children[ci] == null)
    				node.children[ci] = new TrieNode(val);
    			else // update node.sum along the path
    				node.children[ci].sum += val;
    			node = node.children[ci];
    		if (node.isAWord) { // key existed
    			node.isAWord = false;
    			insert(key, -node.val);
    			// for updating all pre nodes'sum along the path
    		node.val = val;
    		node.isAWord = true;
    	public int sum(String prefix) {
    		TrieNode node = root;
    		for (int i = 0; i < prefix.length(); i++) {
    			node = node.children[prefix.charAt(i) - 'a'];
    			if (node == null)
    				return 0;
    		return node.sum;
    class TrieNode {
    	public int val;
    	public int sum;
    	public boolean isAWord;
    	public TrieNode[] children;
    	public TrieNode(int val) {
    		sum = this.val = val;
    		isAWord = false;
    		children = new TrieNode[26];

  • 0

    brilliant idea. Thanks for sharing

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.