LC 677. Map Sum Pairs

Implement a MapSum class with insert, and sum methods.

For the method insert, you'll be given a pair of (string, integer). The string represents the key and the integer represents the value. If the key already existed, then the original key-value pair will be overridden to the new one.

For the method sum, you'll be given a string representing the prefix, and you need to return the sum of all the pairs' value whose key starts with the prefix.

Example 1:

Input: insert("apple", 3), Output: Null
Input: sum("ap"), Output: 3
Input: insert("app", 2), Output: Null
Input: sum("ap"), Output: 5

Runtime: 0 ms, faster than 100.00% of C++ online submissions for Map Sum Pairs.

Trie简单题。

#include <assert.h>
#include <map>
#include <iostream>
#include <vector>
#include <unordered_map>
#include <algorithm>
#define ALL(x) (x).begin(), (x).end()
using namespace std;

struct TrieNode{
  TrieNode* children[26];
  string word;
  int sum;
  TrieNode(){
    for(int i=0; i<26; i++) children[i] = nullptr;
    word = "";
    sum = 0;
  }
};

class Trie{
  TrieNode* root;
public:
  Trie(){
    root = new TrieNode();
  }
  explicit Trie(vector<string> words, vector<int> value){
    root = new TrieNode();
    buildtrie(words, value);
  }
  void buildtrie(vector<string>& words, vector<int>& value){
    for(int i=0; i<words.size(); i++){
      TrieNode* tmp = root;
      for(int j=0; j<words[i].size(); j++){
        int idx = words[i][j] - 'a';
        if(!tmp->children[idx]) tmp->children[idx] = new TrieNode();
        tmp = tmp->children[idx];
        tmp->sum += value[i];
      }
      tmp->word = words[i];
    }
  }
  int getprefixsum(string prefix){
    TrieNode* tmp = root;
    for(int i=0; i<prefix.size(); i++){
      int idx = prefix[i] - 'a';
      if(!tmp->children[idx]) return 0;
      tmp = tmp->children[idx];
    }
    return tmp->sum;
  }
};

class MapSum {
private:
  unordered_map<string, int> mp;
  Trie trie;
public:
  /** Initialize your data structure here. */
  MapSum() {
    trie = Trie();
  }

  void insert(string key, int val) {
    if(!mp.count(key)){
      mp[key] = val;
      vector<string> words = {key};
      vector<int> valvec = {val};
      trie.buildtrie(words,valvec);
    }else{
      vector<string> words = {key};
      vector<int> valvec = {val - mp[key]};
      trie.buildtrie(words, valvec);
    }
  }

  int sum(string prefix) {
    return trie.getprefixsum(prefix);
  }
};