CMU / ts / searchEngine / index.ts

const articles = ["a", "an", "the", "this"];

const prepositions = [
    "in",
    "on",
    "at",
    "to",
    "from",
    "by",
    "with",
    "for",
    "of",
];

const conjunctions = ["and", "or", "but", "yet", "so"];

const pronouns = ["i", "me", "he", "she", "it", "we", "they", "you"];

const auxiliaryVerbs = [
    "is",
    "are",
    "am",
    "be",
    "been",
    "being",
    "has",
    "have",
    "had",
    "do",
    "does",
    "did",
];

const commonVerbs = ["go", "get", "make", "take", "see", "come", "think"];

const adverbs = ["very", "more", "most", "also", "just", "only"];

const otherCommon = [
    "not",
    "no",
    "yes",
    "some",
    "any",
    "all",
    "each",
    "every",
    "what",
    "which",
    "who",
    "when",
    "where",
    "why",
    "how",
];

const contractions = [
    "its",
    "youve",
    "youre",
    "weve",
    "were",
    "itd",
    "youd",
    "yall",
];

const stopWords = new Set([
    ...articles,
    ...prepositions,
    ...adverbs,
    ...otherCommon,
    ...commonVerbs,
    ...conjunctions,
    ...pronouns,
    ...auxiliaryVerbs,
    ...contractions,
]);

export class SearchIndex {
    private index: Map<string, [string, number][]>;

    constructor() {
        this.index = new Map<string, [string, number][]>();
    }

    private getPhrases(
        words: string[],
        filter: boolean[],
    ): Map<string, number> {
        const wordGroups: string[][] = [];
        let currentSlice: string[] = [];
        for (const [index, val] of filter.entries()) {
            if (val) {
                currentSlice.push(words[index]!);
                continue;
            }
            if (currentSlice.length > 1) wordGroups.push(currentSlice);
            currentSlice = [];
        }
        const subPhrases: string[] = wordGroups.flatMap((group) =>
            this.getSubPhrases(group),
        );
        const subPhraseDict = new Map<string, number>();
        for (const sp of subPhrases) {
            subPhraseDict.set(sp, (subPhraseDict.get(sp) || 0) + 1);
        }
        return subPhraseDict;
    }

    private getSubPhrases(phrase: string[]): string[] {
        const subPhrases: string[] = [phrase.join(" ")];
        for (let i = 2; i < phrase.length; i++) {
            for (let offset = 0; offset + i < phrase.length + 1; offset++) {
                const subPhrase = phrase.slice(offset, offset + i).join(" ");
                subPhrases.push(subPhrase);
            }
        }
        return subPhrases;
    }

    private extractKeywords(pageContent: string): Map<string, number> {
        let words: string[] = pageContent
            .split(/\s+/)
            .map((str) => str.replaceAll(/[^\w]+/g, "").toLowerCase())
            .filter((str) => str.length > 0);
        words = [...words, "a"];
        const filter = words.map((word) => !stopWords.has(word));
        const keywords = new Set<string>(words).difference(stopWords);
        const keywordMap = new Map<string, number>(
            keywords.values().map((kw) => [kw, 0]),
        );

        for (let word of words) {
            if (keywords.has(word)) {
                keywordMap.set(word, (keywordMap.get(word) || 0) + 1);
            }
        }
        const phrases = this.getPhrases(words, filter);
        return new Map([...keywordMap, ...phrases]);
    }

    addPage(url: string, pageContent: string): void {
        let keywords = this.extractKeywords(pageContent);
        for (let [kw, count] of keywords.entries()) {
            if (this.index.has(kw)) {
                let prev = this.index.get(kw)!;
                prev.push([url, count]);
                this.index.set(kw, prev);
            } else {
                this.index.set(kw, [[url, count]]);
            }
        }
    }

    updatePage(url: string, pageContent: string): void {
        this.removePage(url);
        this.addPage(url, pageContent);
    }

    removePage(url: string): void {
        this.index.entries().forEach(([keyword, urls]) => {
            const index = urls.findIndex(([u, _]) => u === url);
            if (index >= 0) {
                urls.splice(index, 1);
                if (urls.length === 0) {
                    this.index.delete(keyword);
                }
            }
        });
    }

    checkPage(search: string): boolean {
        for (const urls of this.index.values()) {
            for (const [url, _] of urls) {
                if (search === url) {
                    return true;
                }
            }
        }
        return false;
    }

    size() {
        return this.index.size;
    }

    getPagesForKeyword(keyword: string): string[] {
        const pages = this.index.get(keyword);
        if (!pages) {
            return [];
        }
        return Array.from(pages)
            .sort((a, b) => a[1] - b[1])
            .map(([url, _]) => url);
    }

    search(query: string): string[] {
        const urls = new Map<string, number>();
        const keys = this.extractKeywords(query);
        for (const [key, count] of keys.entries()) {
            const pages = this.index.get(key);
            if (!pages) continue;
            for (const [url, v] of pages) {
                urls.set(
                    url,
                    urls.has(url)
                        ? (urls.get(url) || 0) + v * count
                        : v * count,
                );
            }
        }
        urls.forEach((value, key) => {
            if (key.includes(query)) {
                value += 10;
            }
        });
        return Array.from(urls.entries())
            .sort((a, b) => b[1] - a[1])
            .map(([url, _]) => url);
    }
}