0677.cpp (2128B)
1 // Only Trie 2 class MapSum { 3 struct Node { 4 Node(){}; 5 Node *children[27] = {nullptr}; 6 int &value = reinterpret_cast<int &>(children[0]); 7 int trail = 0; 8 }; 9 10 Node *trie = new Node(); 11 12 public: 13 void insert(const string &key, int val) { 14 Node *prev = trie; 15 for (const char c : key) { 16 const int idx = c & 0x1F; 17 if (!prev->children[idx]) { 18 prev = nullptr; 19 break; 20 }; 21 prev = prev->children[idx]; 22 } 23 24 const int diff = prev ? val - prev->value : val; 25 26 Node *crnt = trie; 27 for (const char c : key) { 28 const int idx = c & 0x1F; 29 crnt->trail += diff; 30 if (!crnt->children[idx]) crnt->children[idx] = new Node(); 31 crnt = crnt->children[idx]; 32 } 33 crnt->value = val; 34 } 35 36 int sum(const string &prefix) const { 37 const Node *crnt = trie; 38 for (const char c : prefix) { 39 const int idx = c & 0x1F; 40 if (!crnt->children[idx]) return 0; 41 crnt = crnt->children[idx]; 42 } 43 return crnt->value + crnt->trail; 44 } 45 }; 46 47 // Trie + hash map 48 class MapSum { 49 struct Node { 50 Node(){}; 51 Node *children[27] = {nullptr}; 52 int &value = reinterpret_cast<int &>(children[0]); 53 }; 54 55 Node *trie = new Node(); 56 unordered_map<string, int> um; 57 58 public: 59 void insert(const string &key, int val) { 60 const int diff = val - um[key]; 61 62 Node *crnt = trie; 63 for (const char c : key) { 64 const int idx = c & 0x1F; 65 crnt->value += diff; 66 if (!crnt->children[idx]) crnt->children[idx] = new Node(); 67 crnt = crnt->children[idx]; 68 } 69 crnt->value += diff; 70 um[key] = val; 71 } 72 73 int sum(const string &prefix) const { 74 const Node *crnt = trie; 75 for (const char c : prefix) { 76 const int idx = c & 0x1F; 77 if (!crnt->children[idx]) return 0; 78 crnt = crnt->children[idx]; 79 } 80 return crnt->value; 81 } 82 };