fork download
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. // =====================================================================
  5. // This file contains a collection of Trie (Prefix Tree) algorithms.
  6. // Each function is ready to be used as a "black box".
  7. // Read the comments above each one to understand:
  8. // - What it solves
  9. // - What input it expects
  10. // - What it returns
  11. // - Time complexity
  12. // - Important constraints / assumptions
  13. //
  14. // TERMINOLOGY:
  15. // - Trie (pronounced "try"): a tree structure used to store strings,
  16. // where each node represents a single character. All descendants of
  17. // a node share a common prefix (hence "prefix tree").
  18. // - Node: an element of the Trie, contains child pointers.
  19. // - Root: the topmost node, represents an empty string.
  20. // - End of word (isEnd): a flag indicating that a complete word ends
  21. // at this node.
  22. // - Prefix count: number of words that have this node as a prefix.
  23. // - Word count: number of times a word has been inserted (useful for
  24. // handling duplicates).
  25. // - Binary Trie (or 0/1 Trie): a Trie over binary strings (bits),
  26. // used for XOR / bitwise operations.
  27. // - XOR (exclusive OR): a bitwise operation, often used with binary
  28. // Tries to find the maximum XOR between numbers.
  29. // =====================================================================
  30.  
  31. // =====================================================================
  32. // 1) Basic Trie for lowercase English letters (a-z)
  33. // This is the most common type used in string problems.
  34. // It supports insert, search, startsWith, and deletion.
  35. // =====================================================================
  36.  
  37. struct TrieNode {
  38. // child[26] for 'a'..'z'
  39. // Using an array is faster and simpler than unordered_map.
  40. TrieNode* child[26];
  41. bool isEnd; // true if a word ends at this node
  42. int wordCount; // number of times this word was inserted
  43. int prefixCount; // number of words that pass through this node
  44.  
  45. TrieNode() {
  46. for (int i = 0; i < 26; i++) child[i] = nullptr;
  47. isEnd = false;
  48. wordCount = 0;
  49. prefixCount = 0;
  50. }
  51. };
  52.  
  53. class Trie {
  54. private:
  55. TrieNode* root;
  56.  
  57. // Helper for deleteWord (recursive)
  58. bool deleteWordHelper(TrieNode* node, const string& word, int depth) {
  59. if (!node) return false;
  60. if (depth == (int)word.size()) {
  61. // We reached the end of the word
  62. if (node->isEnd) {
  63. node->wordCount--;
  64. // If wordCount becomes 0, remove the end marker
  65. if (node->wordCount == 0) {
  66. node->isEnd = false;
  67. }
  68. return true;
  69. }
  70. return false; // word not found
  71. }
  72. int idx = word[depth] - 'a';
  73. bool deleted = deleteWordHelper(node->child[idx], word, depth + 1);
  74. if (deleted) {
  75. // Decrement prefix count of the child
  76. if (node->child[idx]) {
  77. node->child[idx]->prefixCount--;
  78. // If child has no more words and no children, delete it
  79. TrieNode* child = node->child[idx];
  80. if (child->prefixCount == 0 && !child->isEnd) {
  81. bool hasChild = false;
  82. for (int i = 0; i < 26; i++) {
  83. if (child->child[i]) { hasChild = true; break; }
  84. }
  85. if (!hasChild) {
  86. delete child;
  87. node->child[idx] = nullptr;
  88. }
  89. }
  90. }
  91. return true;
  92. }
  93. return false;
  94. }
  95.  
  96. // Helper for getWordsWithPrefix (DFS)
  97. void dfsCollect(TrieNode* node, string& current, vector<string>& result) {
  98. if (node->isEnd) result.push_back(current);
  99. for (int i = 0; i < 26; i++) {
  100. if (node->child[i]) {
  101. current.push_back('a' + i);
  102. dfsCollect(node->child[i], current, result);
  103. current.pop_back();
  104. }
  105. }
  106. }
  107.  
  108. // Helper for smallestWordWithPrefixLength (checks if a path of given length exists)
  109. bool canReachLength(TrieNode* node, int remaining) {
  110. if (remaining == 0) return true;
  111. for (int c = 0; c < 26; c++) {
  112. if (node->child[c] && canReachLength(node->child[c], remaining - 1)) {
  113. return true;
  114. }
  115. }
  116. return false;
  117. }
  118.  
  119. public:
  120. Trie() {
  121. root = new TrieNode();
  122. }
  123.  
  124. // 1.1) Insert a word into the Trie.
  125. // Purpose: Adds the given word to the dictionary.
  126. // Input: word (string) – the word to insert.
  127. // Output: none.
  128. // Time Complexity: O(len(word)), where len is the length of the word.
  129. // Constraints: word contains only lowercase English letters ('a'..'z').
  130. // Note: If the same word is inserted multiple times, each insertion
  131. // increments wordCount and prefixCount for all its nodes.
  132. void insert(const string& word) {
  133. TrieNode* cur = root;
  134. for (char ch : word) {
  135. int idx = ch - 'a';
  136. if (!cur->child[idx]) {
  137. cur->child[idx] = new TrieNode();
  138. }
  139. cur = cur->child[idx];
  140. cur->prefixCount++;
  141. }
  142. cur->isEnd = true;
  143. cur->wordCount++;
  144. }
  145.  
  146. // 1.2) Search for a complete word in the Trie.
  147. // Purpose: Checks if the exact word exists in the Trie.
  148. // Input: word (string) – the word to search.
  149. // Output: returns true if the word exists (at least once), false otherwise.
  150. // Time Complexity: O(len(word)).
  151. // Constraints: word contains only lowercase letters.
  152. // Note: This function does NOT check if the word is a prefix of a longer word;
  153. // it requires isEnd to be true at the last node.
  154. bool search(const string& word) {
  155. TrieNode* cur = root;
  156. for (char ch : word) {
  157. int idx = ch - 'a';
  158. if (!cur->child[idx]) return false;
  159. cur = cur->child[idx];
  160. }
  161. return cur->isEnd;
  162. }
  163.  
  164. // 1.3) Check if there is any word that starts with the given prefix.
  165. // Purpose: Tests whether the given prefix is stored in the Trie.
  166. // Input: prefix (string) – the prefix to check.
  167. // Output: returns true if at least one word has this prefix.
  168. // Time Complexity: O(len(prefix)).
  169. // Constraints: prefix contains only lowercase letters.
  170. // Note: If the prefix itself is a complete word, it still returns true.
  171. bool startsWith(const string& prefix) {
  172. TrieNode* cur = root;
  173. for (char ch : prefix) {
  174. int idx = ch - 'a';
  175. if (!cur->child[idx]) return false;
  176. cur = cur->child[idx];
  177. }
  178. return true;
  179. }
  180.  
  181. // 1.4) Count how many times a given word has been inserted.
  182. // Purpose: Returns the frequency of a specific word (0 if never inserted).
  183. // Input: word (string) – the word to count.
  184. // Output: integer – number of insertions of this exact word.
  185. // Time Complexity: O(len(word)).
  186. // Constraints: word contains only lowercase letters.
  187. int countWord(const string& word) {
  188. TrieNode* cur = root;
  189. for (char ch : word) {
  190. int idx = ch - 'a';
  191. if (!cur->child[idx]) return 0;
  192. cur = cur->child[idx];
  193. }
  194. return cur->wordCount;
  195. }
  196.  
  197. // 1.5) Count how many words have the given prefix.
  198. // Purpose: Returns the total number of inserted words that start
  199. // with the given prefix.
  200. // Input: prefix (string) – the prefix to count.
  201. // Output: integer – number of words with this prefix.
  202. // Time Complexity: O(len(prefix)).
  203. // Constraints: prefix contains only lowercase letters.
  204. int countPrefix(const string& prefix) {
  205. TrieNode* cur = root;
  206. for (char ch : prefix) {
  207. int idx = ch - 'a';
  208. if (!cur->child[idx]) return 0;
  209. cur = cur->child[idx];
  210. }
  211. return cur->prefixCount;
  212. }
  213.  
  214. // 1.6) Delete a word from the Trie (one occurrence).
  215. // Purpose: Removes one occurrence of the given word from the Trie.
  216. // Input: word (string) – the word to delete.
  217. // Output: returns true if the word existed and was deleted,
  218. // false otherwise.
  219. // Time Complexity: O(len(word)).
  220. // Constraints: word contains only lowercase letters.
  221. // Note: This function uses a helper to recursively delete nodes
  222. // that are no longer needed (no children and not end of another word).
  223. bool deleteWord(const string& word) {
  224. return deleteWordHelper(root, word, 0);
  225. }
  226.  
  227. // =====================================================================
  228. // 2) Advanced: Longest Common Prefix among two strings using Trie
  229. // This is a common trick: to find the longest common prefix of
  230. // two strings, you can insert one string and then traverse the other.
  231. // =====================================================================
  232.  
  233. // 2.1) Find the length of the longest common prefix of two strings.
  234. // Purpose: Given two strings a and b, returns the length of their
  235. // longest common prefix.
  236. // Input: a, b (strings) – the two strings to compare.
  237. // Output: integer – length of the longest common prefix.
  238. // Time Complexity: O(len(a) + len(b)) if we insert a first,
  239. // but for a single comparison, it's O(min(len(a), len(b))).
  240. // Constraints: strings contain lowercase letters.
  241. // Note: This implementation inserts a into a temporary Trie,
  242. // then traverses b until mismatch. You can reuse an existing
  243. // Trie if needed.
  244. int longestCommonPrefixLength(const string& a, const string& b) {
  245. // Build a temporary Trie from a
  246. Trie tempTrie;
  247. tempTrie.insert(a);
  248. TrieNode* cur = tempTrie.root;
  249. int len = 0;
  250. for (char ch : b) {
  251. int idx = ch - 'a';
  252. if (!cur->child[idx]) break;
  253. cur = cur->child[idx];
  254. len++;
  255. }
  256. return len;
  257. }
  258.  
  259. // =====================================================================
  260. // 3) Binary Trie (0/1 Trie) for Bitwise XOR problems
  261. // This is a very common pattern in ECPC/ACPC.
  262. // The Trie stores integers as binary strings (bits from MSB to LSB).
  263. // The maximum number of bits is usually 30 or 31 (since numbers up to
  264. // 1e9 require 30 bits, plus sign).
  265. // =====================================================================
  266.  
  267. // Binary Trie node: two children for bit 0 and bit 1.
  268. struct BinaryTrieNode {
  269. BinaryTrieNode* child[2];
  270. int count; // number of numbers that pass through this node
  271. BinaryTrieNode() {
  272. child[0] = child[1] = nullptr;
  273. count = 0;
  274. }
  275. };
  276.  
  277. class BinaryTrie {
  278. private:
  279. BinaryTrieNode* root;
  280. int MAX_BITS; // usually 30 or 31 (e.g., 31 for signed int)
  281.  
  282. // Helper for countLessThanXOR
  283. int countLessThanXORHelper(BinaryTrieNode* node, int x, int limit, int bit) {
  284. if (!node || bit < 0) return 0;
  285. if (limit < 0) return 0;
  286. int xb = (x >> bit) & 1;
  287. int lb = (limit >> bit) & 1;
  288. int res = 0;
  289. if (lb == 1) {
  290. // XOR bit 0 gives a smaller prefix: count all numbers with that bit
  291. if (node->child[xb]) {
  292. res += node->child[xb]->count; // because XOR bit 0 < 1 at this bit
  293. }
  294. // Continue with XOR bit 1 (equal to limit's bit)
  295. if (node->child[xb ^ 1]) {
  296. res += countLessThanXORHelper(node->child[xb ^ 1], x, limit, bit - 1);
  297. }
  298. } else {
  299. // limit bit is 0, so XOR bit must be 0 to stay equal
  300. if (node->child[xb]) {
  301. res += countLessThanXORHelper(node->child[xb], x, limit, bit - 1);
  302. }
  303. }
  304. return res;
  305. }
  306.  
  307. public:
  308. BinaryTrie(int maxBits = 31) {
  309. root = new BinaryTrieNode();
  310. MAX_BITS = maxBits;
  311. }
  312.  
  313. // 3.1) Insert a number into the Binary Trie.
  314. // Purpose: Adds the binary representation of 'num' into the Trie.
  315. // Input: num (int) – the number to insert.
  316. // Output: none.
  317. // Time Complexity: O(MAX_BITS) (constant, ~31 steps).
  318. // Constraints: num fits in a signed 32-bit integer.
  319. // Note: Duplicates are allowed (count is incremented at each node).
  320. void insert(int num) {
  321. BinaryTrieNode* cur = root;
  322. for (int bit = MAX_BITS; bit >= 0; bit--) {
  323. int b = (num >> bit) & 1;
  324. if (!cur->child[b]) cur->child[b] = new BinaryTrieNode();
  325. cur = cur->child[b];
  326. cur->count++;
  327. }
  328. }
  329.  
  330. // 3.2) Query the maximum XOR of 'num' with any number in the Trie.
  331. // Purpose: Given a number, find the maximum XOR value you can get
  332. // by XORing it with any number currently inserted.
  333. // Input: num (int) – the number to XOR against.
  334. // Output: int – the maximum XOR value.
  335. // Time Complexity: O(MAX_BITS).
  336. // Constraints: Trie is not empty.
  337. // Note: This is useful for "Maximum XOR of Two Numbers in an Array".
  338. int maxXOR(int num) {
  339. BinaryTrieNode* cur = root;
  340. int ans = 0;
  341. for (int bit = MAX_BITS; bit >= 0; bit--) {
  342. int b = (num >> bit) & 1;
  343. int desired = b ^ 1; // we want the opposite bit to maximize XOR
  344. if (cur->child[desired]) {
  345. ans |= (1 << bit);
  346. cur = cur->child[desired];
  347. } else {
  348. cur = cur->child[b];
  349. }
  350. }
  351. return ans;
  352. }
  353.  
  354. // 3.3) Query the maximum XOR pair (with a given number) but returns the
  355. // number in the Trie that gives that XOR (not just the value).
  356. // Purpose: Similar to maxXOR, but returns the original number
  357. // from the Trie that yields the max XOR.
  358. // Input: num (int)
  359. // Output: int – the number in the Trie that maximizes XOR.
  360. // Time Complexity: O(MAX_BITS).
  361. // Note: This assumes the Trie has at least one element.
  362. int maxXORNumber(int num) {
  363. BinaryTrieNode* cur = root;
  364. int ans = 0;
  365. for (int bit = MAX_BITS; bit >= 0; bit--) {
  366. int b = (num >> bit) & 1;
  367. int desired = b ^ 1;
  368. if (cur->child[desired]) {
  369. ans |= (desired << bit);
  370. cur = cur->child[desired];
  371. } else {
  372. ans |= (b << bit);
  373. cur = cur->child[b];
  374. }
  375. }
  376. return ans;
  377. }
  378.  
  379. // 3.4) Count numbers in the Trie that are less than a given value
  380. // after XORing with a given number.
  381. // Purpose: For a fixed x, count how many y in the Trie satisfy
  382. // (x ^ y) < limit. This is useful in problems like
  383. // "Count pairs with XOR less than K".
  384. // Input: x (int), limit (int)
  385. // Output: int – count of numbers y such that (x ^ y) < limit.
  386. // Time Complexity: O(MAX_BITS).
  387. // Constraints: Trie is not empty; limit >= 0.
  388. int countLessThanXOR(int x, int limit) {
  389. return countLessThanXORHelper(root, x, limit, MAX_BITS);
  390. }
  391. };
  392.  
  393. // =====================================================================
  394. // 4) Trie with Frequency for Autocomplete / Suggestions
  395. // Not very common in ECPC but can be useful.
  396. // =====================================================================
  397.  
  398. // 4.1) Get all words with a given prefix (lexicographically sorted).
  399. // Purpose: Returns a list of all words that start with the given prefix.
  400. // Input: prefix (string)
  401. // Output: vector<string> – all words in the Trie with that prefix.
  402. // Time Complexity: O(len(prefix) + number of words in the subtree).
  403. // Constraints: words contain lowercase letters.
  404. // Note: This can be memory heavy; use with caution.
  405. vector<string> getWordsWithPrefix(const string& prefix) {
  406. vector<string> result;
  407. TrieNode* cur = root;
  408. for (char ch : prefix) {
  409. int idx = ch - 'a';
  410. if (!cur->child[idx]) return result;
  411. cur = cur->child[idx];
  412. }
  413. string current = prefix;
  414. dfsCollect(cur, current, result);
  415. return result;
  416. }
  417.  
  418. // =====================================================================
  419. // 5) Advanced Idea: Trie + DP for Word Break (not a direct function,
  420. // but a pattern). We include a function that checks if a string
  421. // can be segmented into words from the Trie.
  422. // =====================================================================
  423.  
  424. // 5.1) Word Break: Check if the string can be segmented into words
  425. // that exist in the Trie.
  426. // Purpose: Given a string s, determine if it can be split into
  427. // a sequence of dictionary words (all present in Trie).
  428. // Input: s (string) – the string to segment.
  429. // Output: bool – true if segmentable, false otherwise.
  430. // Time Complexity: O(n^2) in worst case (DP + Trie traversal),
  431. // but can be O(n * maxLen) if we limit.
  432. // Constraints: words in Trie are lowercase; s lowercase.
  433. // Note: This is a classic DP problem. The Trie helps to check
  434. // prefixes quickly.
  435. bool wordBreak(const string& s) {
  436. int n = s.size();
  437. vector<bool> dp(n + 1, false);
  438. dp[0] = true;
  439. for (int i = 0; i < n; i++) {
  440. if (!dp[i]) continue;
  441. TrieNode* cur = root;
  442. for (int j = i; j < n; j++) {
  443. int idx = s[j] - 'a';
  444. if (!cur->child[idx]) break;
  445. cur = cur->child[idx];
  446. if (cur->isEnd) {
  447. dp[j + 1] = true;
  448. }
  449. }
  450. }
  451. return dp[n];
  452. }
  453.  
  454. // =====================================================================
  455. // 6) Trick: Lexicographically smallest string with given prefix and
  456. // length constraints (used in some ECPC problems).
  457. // =====================================================================
  458.  
  459. // 6.1) Find the lexicographically smallest word of length L that
  460. // starts with prefix and is present in the Trie.
  461. // Purpose: Used when you need to construct the smallest string
  462. // that satisfies certain prefix and length constraints.
  463. // Input: prefix (string), length (int)
  464. // Output: string – the smallest word, or empty if none exists.
  465. // Time Complexity: O(length * 26) (DFS over Trie).
  466. // Constraints: Trie must contain at least one such word.
  467. string smallestWordWithPrefixLength(const string& prefix, int length) {
  468. TrieNode* cur = root;
  469. string res = prefix;
  470. // First, traverse to the end of the prefix
  471. for (char ch : prefix) {
  472. int idx = ch - 'a';
  473. if (!cur->child[idx]) return "";
  474. cur = cur->child[idx];
  475. }
  476. if ((int)res.size() > length) return ""; // prefix already longer
  477. // Greedily choose the smallest character that leads to a valid word
  478. while ((int)res.size() < length) {
  479. bool found = false;
  480. for (int c = 0; c < 26; c++) {
  481. if (cur->child[c]) {
  482. // Check if this branch can reach the required length
  483. if (canReachLength(cur->child[c], length - (int)res.size() - 1)) {
  484. res.push_back('a' + c);
  485. cur = cur->child[c];
  486. found = true;
  487. break;
  488. }
  489. }
  490. }
  491. if (!found) return "";
  492. }
  493. return res;
  494. }
  495.  
  496. // =====================================================================
  497. // 7) Aho-Corasick (Advanced) – not fully implemented, just a note.
  498. // For multiple pattern matching, you can build a trie with failure
  499. // links. This is a common advanced topic in ECPC/ACPC.
  500. // If you need it, implement a separate class for Aho-Corasick.
  501. // It is essentially a Trie with BFS failure links.
  502. // =====================================================================
  503.  
  504. // =====================================================================
  505. // 8) Important Notes / Warnings
  506. // - Always ensure the character set is consistent (e.g., lowercase).
  507. // - Memory: each node has 26 pointers; for large data, it can be
  508. // heavy. Use unordered_map or vector of ints (compressed trie)
  509. // if memory is an issue.
  510. // - For Binary Trie, MAX_BITS depends on the maximum number you
  511. // will insert. Usually 30 for up to 1e9, 31 if including sign.
  512. // - The wordBreak and getWordsWithPrefix functions are provided as
  513. // examples of advanced usage; adjust as needed.
  514. // - Deletion is tricky; our deleteWord handles one occurrence, but
  515. // you may need to adjust for your use case.
  516. // =====================================================================
  517. };
  518.  
  519. // =====================================================================
  520. // Example usage (you can remove this when using as a black box)
  521. // =====================================================================
  522. int main() {
  523. ios::sync_with_stdio(false);
  524. cin.tie(nullptr);
  525.  
  526. // Basic Trie
  527. Trie trie;
  528. trie.insert("apple");
  529. trie.insert("app");
  530. trie.insert("apricot");
  531. cout << boolalpha;
  532. cout << "search('app'): " << trie.search("app") << '\n'; // true
  533. cout << "search('apple'): " << trie.search("apple") << '\n'; // true
  534. cout << "search('ap'): " << trie.search("ap") << '\n'; // false
  535. cout << "startsWith('ap'): " << trie.startsWith("ap") << '\n'; // true
  536. cout << "countPrefix('ap'): " << trie.countPrefix("ap") << '\n'; // 3 (apple, app, apricot)
  537. cout << "countWord('app'): " << trie.countWord("app") << '\n'; // 1
  538.  
  539. // Binary Trie for XOR
  540. Trie::BinaryTrie bt(30);
  541. bt.insert(5); // 101
  542. bt.insert(2); // 010
  543. bt.insert(7); // 111
  544. cout << "maxXOR with 1: " << bt.maxXOR(1) << '\n'; // 6 (1 xor 7 = 6)
  545. cout << "maxXOR with 3: " << bt.maxXOR(3) << '\n'; // 4 (3 xor 7 = 4)
  546.  
  547. // Word Break
  548. Trie dict;
  549. dict.insert("leet");
  550. dict.insert("code");
  551. cout << "wordBreak('leetcode'): " << dict.wordBreak("leetcode") << '\n'; // true
  552. cout << "wordBreak('leetocode'): " << dict.wordBreak("leetocode") << '\n'; // false
  553.  
  554. // Longest common prefix
  555. cout << "LCP('abcdef', 'abcxyz'): " << trie.longestCommonPrefixLength("abcdef", "abcxyz") << '\n'; // 3
  556.  
  557. return 0;
  558. }
Success #stdin #stdout 0s 5324KB
stdin
Standard input is empty
stdout
search('app'): true
search('apple'): true
search('ap'): false
startsWith('ap'): true
countPrefix('ap'): 3
countWord('app'): 1
maxXOR with 1: 6
maxXOR with 3: 6
wordBreak('leetcode'): true
wordBreak('leetocode'): false
LCP('abcdef', 'abcxyz'): 3