Refactor: further improve trie implementation

This commit is contained in:
SukkaW 2023-12-30 19:12:11 +08:00
parent ef6aa0f8ab
commit a7c1b0fd17

View File

@ -5,36 +5,35 @@
export const SENTINEL = '\u0000'; export const SENTINEL = '\u0000';
type TrieNode = { type TrieNode = {
[SENTINEL]?: true [SENTINEL]: boolean
} & { } & {
[key: string & {}]: TrieNode | undefined [key: string & {}]: TrieNode | undefined
}; };
// type TrieNode = Map<typeof SENTINEL | string & {}, TrieNode | true | undefined>; const createNode = (): TrieNode => ({
[SENTINEL]: false
}) as TrieNode;
/**
* @param {string[] | Set<string>} [from]
*/
export const createTrie = (from?: string[] | Set<string>) => { export const createTrie = (from?: string[] | Set<string>) => {
let size = 0; let size = 0;
const root: TrieNode = {}; const root: TrieNode = createNode();
/** /**
* Method used to add the given prefix to the trie. * Method used to add the given prefix to the trie.
*
* @param {string} suffix - Prefix to follow.
*/ */
const add = (suffix: string): void => { const add = (suffix: string): void => {
let node: TrieNode = root; let node: TrieNode = root;
let token: string; let token: string;
for (let i = suffix.length - 1; i >= 0; i--) { for (let i = suffix.length - 1; i >= 0; i--) {
token = suffix[i]; token = suffix[i];
node[token] ||= {}; if (!(token in node)) {
node = node[token]!; node[token] = createNode();
}
node = node[token]!; // we know it is defined
} }
// Do we need to increase size? // Do we need to increase size?
if (!(SENTINEL in node)) { if (!node[SENTINEL]) {
size++; size++;
} }
node[SENTINEL] = true; node[SENTINEL] = true;
@ -50,21 +49,14 @@ export const createTrie = (from?: string[] | Set<string>) => {
for (let i = suffix.length - 1; i >= 0; i--) { for (let i = suffix.length - 1; i >= 0; i--) {
token = suffix[i]; token = suffix[i];
const n = node[token]; if (!(token in node)) return false;
if (n === undefined) return false; node = node[token]!; // we know it is defined
// if (n === true) return false;
node = n;
} }
return true; return true;
}; };
/** /**
* Method used to retrieve every item in the trie with the given prefix. * Method used to retrieve every item in the trie with the given prefix.
*
* @param {string} suffix - Prefix to query.
* @param {boolean} [includeEqualWithSuffix]
* @return {string[]}
*/ */
const find = (suffix: string, includeEqualWithSuffix = true): string[] => { const find = (suffix: string, includeEqualWithSuffix = true): string[] => {
let node: TrieNode = root; let node: TrieNode = root;
@ -73,11 +65,8 @@ export const createTrie = (from?: string[] | Set<string>) => {
for (let i = suffix.length - 1; i >= 0; i--) { for (let i = suffix.length - 1; i >= 0; i--) {
token = suffix[i]; token = suffix[i];
const n = node[token]; if (!(token in node)) return [];
if (n === undefined) return []; node = node[token]!;
// if (n === true) return [];
node = n;
} }
const matches: string[] = []; const matches: string[] = [];
@ -94,16 +83,14 @@ export const createTrie = (from?: string[] | Set<string>) => {
$suffix = suffixStack.pop()!; $suffix = suffixStack.pop()!;
node = nodeStack.pop()!; node = nodeStack.pop()!;
// eslint-disable-next-line guard-for-in -- plain object if (node[SENTINEL]) {
for (k in node) { if (includeEqualWithSuffix || $suffix !== suffix) {
if (k === SENTINEL) { matches.push($suffix);
if (includeEqualWithSuffix || $suffix !== suffix) {
matches.push($suffix);
}
continue;
} }
}
for (k in node) {
if (k === SENTINEL) continue;
nodeStack.push(node[k]!); nodeStack.push(node[k]!);
suffixStack.push(k + $suffix); suffixStack.push(k + $suffix);
} }
@ -114,9 +101,6 @@ export const createTrie = (from?: string[] | Set<string>) => {
/** /**
* Method used to delete a prefix from the trie. * Method used to delete a prefix from the trie.
*
* @param {string} suffix - Prefix to delete.
* @return {boolean}
*/ */
const remove = (suffix: string): boolean => { const remove = (suffix: string): boolean => {
let node: TrieNode = root; let node: TrieNode = root;
@ -129,12 +113,10 @@ export const createTrie = (from?: string[] | Set<string>) => {
token = suffix[i]; token = suffix[i];
parent = node; parent = node;
const n = node[token];
// Prefix does not exist // Prefix does not exist
if (n === undefined) return false; if (!(token in node)) return false;
// if (n === true) return false // if (n === true) return false
node = node[token]!; // we know it is defined
node = n;
// Keeping track of a potential branch to prune // Keeping track of a potential branch to prune
// If the node is to be pruned, but they are more than one token child in it, we can't prune it // If the node is to be pruned, but they are more than one token child in it, we can't prune it
@ -159,14 +141,14 @@ export const createTrie = (from?: string[] | Set<string>) => {
} }
} }
if (!(SENTINEL in node)) return false; if (!node[SENTINEL]) return false;
size--; size--;
if (tokenToPrune && toPrune) { if (tokenToPrune && toPrune) {
delete toPrune[tokenToPrune]; delete toPrune[tokenToPrune];
} else { } else {
delete node[SENTINEL]; node[SENTINEL] = false;
} }
return true; return true;
@ -174,24 +156,17 @@ export const createTrie = (from?: string[] | Set<string>) => {
/** /**
* Method used to assert whether the given prefix exists in the Trie. * Method used to assert whether the given prefix exists in the Trie.
*
* @param {string} suffix - Prefix to check.
* @return {boolean}
*/ */
const has = (suffix: string): boolean => { const has = (suffix: string): boolean => {
let node: TrieNode = root; let node: TrieNode = root;
for (let i = suffix.length - 1; i >= 0; i--) { for (let i = suffix.length - 1; i >= 0; i--) {
const n = node[suffix[i]]; const token = suffix[i];
if (n === undefined) { if (!(token in node)) return false;
return false; node = node[token]!; // we know it is defined
}
// if (n === true) return false;
node = n;
} }
return SENTINEL in node; return node[SENTINEL];
}; };
if (from) { if (from) {
@ -207,6 +182,9 @@ export const createTrie = (from?: string[] | Set<string>) => {
has, has,
get size() { get size() {
return size; return size;
},
get root() {
return root;
} }
}; };
}; };