Skip to content

Commit 732f7c8

Browse files
authored
Add BM25 Inverted Index Search Algorithm (#5615)
1 parent 676d451 commit 732f7c8

File tree

2 files changed

+313
-0
lines changed

2 files changed

+313
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
package com.thealgorithms.searches;
2+
3+
import java.util.ArrayList;
4+
import java.util.HashMap;
5+
import java.util.List;
6+
import java.util.Map;
7+
import java.util.Objects;
8+
9+
/**
10+
* Inverted Index implementation with BM25 Scoring for movie search.
11+
* This class supports adding movie documents and searching for terms
12+
* within those documents using the BM25 algorithm.
13+
* @author Prayas Kumar (https://github.com/prayas7102)
14+
*/
15+
16+
class Movie {
17+
int docId; // Unique identifier for the movie
18+
String name; // Movie name
19+
double imdbRating; // IMDb rating of the movie
20+
int releaseYear; // Year the movie was released
21+
String content; // Full text content (could be the description or script)
22+
23+
/**
24+
* Constructor for the Movie class.
25+
* @param docId Unique identifier for the movie.
26+
* @param name Name of the movie.
27+
* @param imdbRating IMDb rating of the movie.
28+
* @param releaseYear Release year of the movie.
29+
* @param content Content or description of the movie.
30+
*/
31+
Movie(int docId, String name, double imdbRating, int releaseYear, String content) {
32+
this.docId = docId;
33+
this.name = name;
34+
this.imdbRating = imdbRating;
35+
this.releaseYear = releaseYear;
36+
this.content = content;
37+
}
38+
39+
/**
40+
* Get all the words from the movie's name and content.
41+
* Converts the name and content to lowercase and splits on non-word characters.
42+
* @return Array of words from the movie name and content.
43+
*/
44+
public String[] getWords() {
45+
return (name + " " + content).toLowerCase().split("\\W+");
46+
}
47+
48+
@Override
49+
public String toString() {
50+
return "Movie{"
51+
+ "docId=" + docId + ", name='" + name + '\'' + ", imdbRating=" + imdbRating + ", releaseYear=" + releaseYear + '}';
52+
}
53+
}
54+
55+
class SearchResult {
56+
int docId; // Unique identifier of the movie document
57+
double relevanceScore; // Relevance score based on the BM25 algorithm
58+
59+
/**
60+
* Constructor for SearchResult class.
61+
* @param docId Document ID (movie) for this search result.
62+
* @param relevanceScore The relevance score based on BM25 scoring.
63+
*/
64+
SearchResult(int docId, double relevanceScore) {
65+
this.docId = docId;
66+
this.relevanceScore = relevanceScore;
67+
}
68+
69+
public int getDocId() {
70+
return docId;
71+
}
72+
73+
@Override
74+
public String toString() {
75+
return "SearchResult{"
76+
+ "docId=" + docId + ", relevanceScore=" + relevanceScore + '}';
77+
}
78+
79+
@Override
80+
public boolean equals(Object o) {
81+
if (this == o) {
82+
return true;
83+
}
84+
if (o == null || getClass() != o.getClass()) {
85+
return false;
86+
}
87+
SearchResult that = (SearchResult) o;
88+
return docId == that.docId && Double.compare(that.relevanceScore, relevanceScore) == 0;
89+
}
90+
91+
@Override
92+
public int hashCode() {
93+
return Objects.hash(docId, relevanceScore);
94+
}
95+
96+
public double getRelevanceScore() {
97+
return this.relevanceScore;
98+
}
99+
}
100+
101+
public final class BM25InvertedIndex {
102+
private Map<String, Map<Integer, Integer>> index; // Inverted index mapping terms to document id and frequency
103+
private Map<Integer, Movie> movies; // Mapping of movie document IDs to Movie objects
104+
private int totalDocuments; // Total number of movies/documents
105+
private double avgDocumentLength; // Average length of documents (number of words)
106+
private static final double K = 1.5; // BM25 tuning parameter, controls term frequency saturation
107+
private static final double B = 0.75; // BM25 tuning parameter, controls length normalization
108+
109+
/**
110+
* Constructor for BM25InvertedIndex.
111+
* Initializes the inverted index and movie storage.
112+
*/
113+
BM25InvertedIndex() {
114+
index = new HashMap<>();
115+
movies = new HashMap<>();
116+
totalDocuments = 0;
117+
avgDocumentLength = 0.0;
118+
}
119+
120+
/**
121+
* Add a movie to the index.
122+
* @param docId Unique identifier for the movie.
123+
* @param name Name of the movie.
124+
* @param imdbRating IMDb rating of the movie.
125+
* @param releaseYear Release year of the movie.
126+
* @param content Content or description of the movie.
127+
*/
128+
public void addMovie(int docId, String name, double imdbRating, int releaseYear, String content) {
129+
Movie movie = new Movie(docId, name, imdbRating, releaseYear, content);
130+
movies.put(docId, movie);
131+
totalDocuments++;
132+
133+
// Get words (terms) from the movie's name and content
134+
String[] terms = movie.getWords();
135+
int docLength = terms.length;
136+
137+
// Update the average document length
138+
avgDocumentLength = (avgDocumentLength * (totalDocuments - 1) + docLength) / totalDocuments;
139+
140+
// Update the inverted index
141+
for (String term : terms) {
142+
// Create a new entry if the term is not yet in the index
143+
index.putIfAbsent(term, new HashMap<>());
144+
145+
// Get the list of documents containing the term
146+
Map<Integer, Integer> docList = index.get(term);
147+
if (docList == null) {
148+
docList = new HashMap<>();
149+
index.put(term, docList); // Ensure docList is added to the index
150+
}
151+
// Increment the term frequency in this document
152+
docList.put(docId, docList.getOrDefault(docId, 0) + 1);
153+
}
154+
}
155+
156+
public int getMoviesLength() {
157+
return movies.size();
158+
}
159+
160+
/**
161+
* Search for documents containing a term using BM25 scoring.
162+
* @param term The search term.
163+
* @return A list of search results sorted by relevance score.
164+
*/
165+
public List<SearchResult> search(String term) {
166+
term = term.toLowerCase(); // Normalize search term
167+
if (!index.containsKey(term)) {
168+
return new ArrayList<>(); // Return empty list if term not found
169+
}
170+
171+
Map<Integer, Integer> termDocs = index.get(term); // Documents containing the term
172+
List<SearchResult> results = new ArrayList<>();
173+
174+
// Compute IDF for the search term
175+
double idf = computeIDF(termDocs.size());
176+
177+
// Calculate relevance scores for all documents containing the term
178+
for (Map.Entry<Integer, Integer> entry : termDocs.entrySet()) {
179+
int docId = entry.getKey();
180+
int termFrequency = entry.getValue();
181+
Movie movie = movies.get(docId);
182+
if (movie == null) {
183+
continue; // Skip this document if movie doesn't exist
184+
}
185+
double docLength = movie.getWords().length;
186+
187+
// Compute BM25 relevance score
188+
double score = computeBM25Score(termFrequency, docLength, idf);
189+
results.add(new SearchResult(docId, score));
190+
}
191+
192+
// Sort the results by relevance score in descending order
193+
results.sort((r1, r2) -> Double.compare(r2.relevanceScore, r1.relevanceScore));
194+
return results;
195+
}
196+
197+
/**
198+
* Compute the BM25 score for a given term and document.
199+
* @param termFrequency The frequency of the term in the document.
200+
* @param docLength The length of the document.
201+
* @param idf The inverse document frequency of the term.
202+
* @return The BM25 relevance score for the term in the document.
203+
*/
204+
private double computeBM25Score(int termFrequency, double docLength, double idf) {
205+
double numerator = termFrequency * (K + 1);
206+
double denominator = termFrequency + K * (1 - B + B * (docLength / avgDocumentLength));
207+
return idf * (numerator / denominator);
208+
}
209+
210+
/**
211+
* Compute the inverse document frequency (IDF) of a term.
212+
* The IDF measures the importance of a term across the entire document set.
213+
* @param docFrequency The number of documents that contain the term.
214+
* @return The inverse document frequency (IDF) value.
215+
*/
216+
private double computeIDF(int docFrequency) {
217+
// Total number of documents in the index
218+
return Math.log((totalDocuments - docFrequency + 0.5) / (docFrequency + 0.5));
219+
}
220+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package com.thealgorithms.searches;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertFalse;
5+
import static org.junit.jupiter.api.Assertions.assertTrue;
6+
7+
import java.util.List;
8+
import org.junit.jupiter.api.BeforeAll;
9+
import org.junit.jupiter.api.Test;
10+
11+
/**
12+
* Test Cases for Inverted Index with BM25
13+
* @author Prayas Kumar (https://github.com/prayas7102)
14+
*/
15+
16+
class BM25InvertedIndexTest {
17+
18+
private static BM25InvertedIndex index;
19+
20+
@BeforeAll
21+
static void setUp() {
22+
index = new BM25InvertedIndex();
23+
index.addMovie(1, "The Shawshank Redemption", 9.3, 1994, "Hope is a good thing. Maybe the best of things. And no good thing ever dies.");
24+
index.addMovie(2, "The Godfather", 9.2, 1972, "I'm gonna make him an offer he can't refuse.");
25+
index.addMovie(3, "The Dark Knight", 9.0, 2008, "You either die a hero or live long enough to see yourself become the villain.");
26+
index.addMovie(4, "Pulp Fiction", 8.9, 1994, "You know what they call a Quarter Pounder with Cheese in Paris? They call it a Royale with Cheese.");
27+
index.addMovie(5, "Good Will Hunting", 8.3, 1997, "Will Hunting is a genius and he has a good heart. The best of his abilities is yet to be explored.");
28+
index.addMovie(6, "It's a Wonderful Life", 8.6, 1946, "Each man's life touches so many other lives. If he wasn't around, it would leave an awfully good hole.");
29+
index.addMovie(7, "The Pursuit of Happyness", 8.0, 2006, "It was the pursuit of a better life, and a good opportunity to change things for the better.");
30+
index.addMovie(8, "A Few Good Men", 7.7, 1992, "You can't handle the truth! This movie has a lot of good moments and intense drama.");
31+
}
32+
33+
@Test
34+
void testAddMovie() {
35+
// Check that the index contains the correct number of movies
36+
int moviesLength = index.getMoviesLength();
37+
assertEquals(8, moviesLength);
38+
}
39+
40+
@Test
41+
void testSearchForTermFound() {
42+
int expected = 1;
43+
List<SearchResult> result = index.search("hope");
44+
int actual = result.getFirst().getDocId();
45+
assertEquals(expected, actual);
46+
}
47+
48+
@Test
49+
void testSearchRanking() {
50+
// Perform search for the term "good"
51+
List<SearchResult> results = index.search("good");
52+
assertFalse(results.isEmpty());
53+
54+
// Validate the ranking based on the provided relevance scores
55+
assertEquals(6, results.get(0).getDocId()); // It's a Wonderful Life should be ranked 1st
56+
assertEquals(7, results.get(1).getDocId()); // The Pursuit of Happyness should be ranked 2nd
57+
assertEquals(5, results.get(2).getDocId()); // Good Will Hunting should be ranked 3rd
58+
assertEquals(8, results.get(3).getDocId()); // A Few Good Men should be ranked 4th
59+
assertEquals(1, results.get(4).getDocId()); // The Shawshank Redemption should be ranked 5th
60+
61+
// Ensure the relevance scores are in descending order
62+
for (int i = 0; i < results.size() - 1; i++) {
63+
assertTrue(results.get(i).getRelevanceScore() > results.get(i + 1).getRelevanceScore());
64+
}
65+
}
66+
67+
@Test
68+
void testSearchForTermNotFound() {
69+
List<SearchResult> results = index.search("nonexistent");
70+
assertTrue(results.isEmpty());
71+
}
72+
73+
@Test
74+
void testSearchForCommonTerm() {
75+
List<SearchResult> results = index.search("the");
76+
assertFalse(results.isEmpty());
77+
assertTrue(results.size() > 1);
78+
}
79+
80+
@Test
81+
void testBM25ScoreCalculation() {
82+
List<SearchResult> results = index.search("cheese");
83+
assertEquals(1, results.size());
84+
assertEquals(4, results.getFirst().docId); // Pulp Fiction should have the highest score
85+
}
86+
87+
@Test
88+
void testCaseInsensitivity() {
89+
List<SearchResult> resultsLowerCase = index.search("hope");
90+
List<SearchResult> resultsUpperCase = index.search("HOPE");
91+
assertEquals(resultsLowerCase, resultsUpperCase);
92+
}
93+
}

0 commit comments

Comments
 (0)