InvertedIndex.java

import java.util.*;

/**
 * InvertedIndex maps:
 *   word -> (documentId -> frequency in that document)
 */
public class InvertedIndex {

    private final Map<String, Map<Integer, Integer>> index;
    private int totalDocuments;

    public InvertedIndex() {
        index = new HashMap<>();
        totalDocuments = 0;
    }

    /**
     * Adds a document to the index.
     */
    public void addDocument(Document document) {

        totalDocuments++;

        List<String> words = Tokenizer.tokenize(document.getContent());

        for (String word : words) {
            index.putIfAbsent(word, new HashMap<>());

            Map<Integer, Integer> docFrequencyMap = index.get(word);

            int docId = document.getId();
            docFrequencyMap.put(
                docId,
                docFrequencyMap.getOrDefault(docId, 0) + 1
            );
        }
    }

    /**
     * Returns map of documentId -> frequency for a given word.
     */
    public Map<Integer, Integer> getDocuments(String word) {
        return index.getOrDefault(word, Collections.emptyMap());
    }

    /**
     * Number of documents indexed.
     */
    public int getTotalDocuments() {
        return totalDocuments;
    }

    /**
     * Document Frequency (DF) of a word.
     */
    public int getDocumentFrequency(String word) {
        return index.containsKey(word) ? index.get(word).size() : 0;
    }

    /**
     * Check if word exists in index.
     */
    public boolean containsWord(String word) {
        return index.containsKey(word);
    }
}