import string
from collections import defaultdict


# Step 1: Preprocess text
def preprocess_text(text):
    text = text.lower()
    text = text.translate(str.maketrans("", "", string.punctuation))
    return text.split()


# Step 2: Build inverted index
def build_inverted_index(documents):
    inverted_index = defaultdict(set)

    for doc_id, text in documents.items():
        words = preprocess_text(text)
        for word in words:
            inverted_index[word].add(doc_id)

    return inverted_index


# Step 3: Search function (AND query)
def search(inverted_index, query):
    query_terms = preprocess_text(query)
    result_set = None

    for term in query_terms:
        if term in inverted_index:
            if result_set is None:
                result_set = inverted_index[term]
            else:
                result_set = result_set.intersection(inverted_index[term])
        else:
            return set()

    return result_set if result_set else set()


# Sample documents
documents = {
    1: "Information retrieval is an essential aspect of search engines.",
    2: "The field of information retrieval focuses on algorithm.",
    3: "Search engines use retrieval techniques to improve performance.",
    4: "Deep learning models are used for information retrieval tasks."
}


# Build index
inverted_index = build_inverted_index(documents)


# Print inverted index
print("Inverted Index (Alphabetical Order):")
for term in sorted(inverted_index.keys()):
    print(f"{term} : {sorted(inverted_index[term])}")


# Search query
query = "retrieval"
result = search(inverted_index, query)

print(f"\nDocuments containing the query '{query}': {sorted(result)}")