BM25InvertedIndex.java

packagecom.thealgorithms.searches;

importjava.util.ArrayList;
importjava.util.HashMap;
importjava.util.List;
importjava.util.Map;
importjava.util.Objects;

/**
 * Inverted Index implementation with BM25 Scoring for movie search.
 * This class supports adding movie documents and searching for terms
 * within those documents using the BM25 algorithm.
 * @author Prayas Kumar (https://github.com/prayas7102)
 */

classMovie {
intdocId; // Unique identifier for the movie
Stringname; // Movie name
doubleimdbRating; // IMDb rating of the movie
intreleaseYear; // Year the movie was released
Stringcontent; // Full text content (could be the description or script)

/**
 * Constructor for the Movie class.
 * @param docId Unique identifier for the movie.
 * @param name Name of the movie.
 * @param imdbRating IMDb rating of the movie.
 * @param releaseYear Release year of the movie.
 * @param content Content or description of the movie.
 */
Movie(intdocId, Stringname, doubleimdbRating, intreleaseYear, Stringcontent) {
this.docId = docId;
this.name = name;
this.imdbRating = imdbRating;
this.releaseYear = releaseYear;
this.content = content;
 }

/**
 * Get all the words from the movie's name and content.
 * Converts the name and content to lowercase and splits on non-word characters.
 * @return Array of words from the movie name and content.
 */
publicString[] getWords() {
return (name + " " + content).toLowerCase().split("\\W+");
 }

@Override
publicStringtoString() {
return"Movie{"
 + "docId=" + docId + ", name='" + name + '\'' + ", imdbRating=" + imdbRating + ", releaseYear=" + releaseYear + '}';
 }
}

classSearchResult {
intdocId; // Unique identifier of the movie document
doublerelevanceScore; // Relevance score based on the BM25 algorithm

/**
 * Constructor for SearchResult class.
 * @param docId Document ID (movie) for this search result.
 * @param relevanceScore The relevance score based on BM25 scoring.
 */
SearchResult(intdocId, doublerelevanceScore) {
this.docId = docId;
this.relevanceScore = relevanceScore;
 }

publicintgetDocId() {
returndocId;
 }

@Override
publicStringtoString() {
return"SearchResult{"
 + "docId=" + docId + ", relevanceScore=" + relevanceScore + '}';
 }

@Override
publicbooleanequals(Objecto) {
if (this == o) {
returntrue;
 }
if (o == null || getClass() != o.getClass()) {
returnfalse;
 }
SearchResultthat = (SearchResult) o;
returndocId == that.docId && Double.compare(that.relevanceScore, relevanceScore) == 0;
 }

@Override
publicinthashCode() {
returnObjects.hash(docId, relevanceScore);
 }

publicdoublegetRelevanceScore() {
returnthis.relevanceScore;
 }
}

publicfinalclassBM25InvertedIndex {
privateMap<String, Map<Integer, Integer>> index; // Inverted index mapping terms to document id and frequency
privateMap<Integer, Movie> movies; // Mapping of movie document IDs to Movie objects
privateinttotalDocuments; // Total number of movies/documents
privatedoubleavgDocumentLength; // Average length of documents (number of words)
privatestaticfinaldoubleK = 1.5; // BM25 tuning parameter, controls term frequency saturation
privatestaticfinaldoubleB = 0.75; // BM25 tuning parameter, controls length normalization

/**
 * Constructor for BM25InvertedIndex.
 * Initializes the inverted index and movie storage.
 */
BM25InvertedIndex() {
index = newHashMap<>();
movies = newHashMap<>();
totalDocuments = 0;
avgDocumentLength = 0.0;
 }

/**
 * Add a movie to the index.
 * @param docId Unique identifier for the movie.
 * @param name Name of the movie.
 * @param imdbRating IMDb rating of the movie.
 * @param releaseYear Release year of the movie.
 * @param content Content or description of the movie.
 */
publicvoidaddMovie(intdocId, Stringname, doubleimdbRating, intreleaseYear, Stringcontent) {
Moviemovie = newMovie(docId, name, imdbRating, releaseYear, content);
movies.put(docId, movie);
totalDocuments++;

// Get words (terms) from the movie's name and content
String[] terms = movie.getWords();
intdocLength = terms.length;

// Update the average document length
avgDocumentLength = (avgDocumentLength * (totalDocuments - 1) + docLength) / totalDocuments;

// Update the inverted index
for (Stringterm : terms) {
// Create a new entry if the term is not yet in the index
index.putIfAbsent(term, newHashMap<>());

// Get the list of documents containing the term
Map<Integer, Integer> docList = index.get(term);
if (docList == null) {
docList = newHashMap<>();
index.put(term, docList); // Ensure docList is added to the index
 }
// Increment the term frequency in this document
docList.put(docId, docList.getOrDefault(docId, 0) + 1);
 }
 }

publicintgetMoviesLength() {
returnmovies.size();
 }

/**
 * Search for documents containing a term using BM25 scoring.
 * @param term The search term.
 * @return A list of search results sorted by relevance score.
 */
publicList<SearchResult> search(Stringterm) {
term = term.toLowerCase(); // Normalize search term
if (!index.containsKey(term)) {
returnnewArrayList<>(); // Return empty list if term not found
 }

Map<Integer, Integer> termDocs = index.get(term); // Documents containing the term
List<SearchResult> results = newArrayList<>();

// Compute IDF for the search term
doubleidf = computeIDF(termDocs.size());

// Calculate relevance scores for all documents containing the term
for (Map.Entry<Integer, Integer> entry : termDocs.entrySet()) {
intdocId = entry.getKey();
inttermFrequency = entry.getValue();
Moviemovie = movies.get(docId);
if (movie == null) {
continue; // Skip this document if movie doesn't exist
 }
doubledocLength = movie.getWords().length;

// Compute BM25 relevance score
doublescore = computeBM25Score(termFrequency, docLength, idf);
results.add(newSearchResult(docId, score));
 }

// Sort the results by relevance score in descending order
results.sort((r1, r2) -> Double.compare(r2.relevanceScore, r1.relevanceScore));
returnresults;
 }

/**
 * Compute the BM25 score for a given term and document.
 * @param termFrequency The frequency of the term in the document.
 * @param docLength The length of the document.
 * @param idf The inverse document frequency of the term.
 * @return The BM25 relevance score for the term in the document.
 */
privatedoublecomputeBM25Score(inttermFrequency, doubledocLength, doubleidf) {
doublenumerator = termFrequency * (K + 1);
doubledenominator = termFrequency + K * (1 - B + B * (docLength / avgDocumentLength));
returnidf * (numerator / denominator);
 }

/**
 * Compute the inverse document frequency (IDF) of a term.
 * The IDF measures the importance of a term across the entire document set.
 * @param docFrequency The number of documents that contain the term.
 * @return The inverse document frequency (IDF) value.
 */
privatedoublecomputeIDF(intdocFrequency) {
// Total number of documents in the index
returnMath.log((totalDocuments - docFrequency + 0.5) / (docFrequency + 0.5) + 1);
 }
}