1 files changed, 34 insertions, 20 deletions

diff --git a/search/search.go b/search/search.go
index f0b1cce..5fb96f3 100644
--- a/search/search.go
+++ b/search/search.go
@@ -3,11 +3,16 @@
 import (
 	"searchEngine/analyzer"
 	"searchEngine/index"
+	"slices"
 )
 
+type Result struct {
+	DocID int
+	Score int
+}
+
 func Search(idx *index.InvertedIndex, query string) []string {
 	tokenizedQuery := analyzer.ProcessText(query)
-
 	if len(tokenizedQuery) == 0 {
 		return nil
 	}
@@ -16,43 +21,52 @@ func Search(idx *index.InvertedIndex, query string) []string {
 	defer idx.Mu.RUnlock()
 
 	firstToken := tokenizedQuery[0]
-	baseIDs, exists := idx.Data[firstToken]
+	basePostings, exists := idx.Data[firstToken]
 	if !exists {
 		return nil
 	}
 
-	intersection := make([]index.Posting, len(baseIDs))
-	copy(intersection, baseIDs)
+	scores := make(map[int]int)
+	for _, p := range basePostings {
+		scores[p.DocId] = p.Count
+	}
 
 	for i := 1; i < len(tokenizedQuery); i++ {
 		token := tokenizedQuery[i]
-
-		nextIDs, exists := idx.Data[token]
+		nextPostings, exists := idx.Data[token]
 		if !exists {
 			return nil
 		}
 
-		lookup := make(map[int]struct{})
-		for _, id := range nextIDs {
-			lookup[id.DocId] = struct{}{}
+		lookup := make(map[int]int)
+		for _, p := range nextPostings {
+			lookup[p.DocId] = p.Count
 		}
 
-		var filtered []index.Posting
-		for _, id := range intersection {
-			if _, found := lookup[id.DocId]; found {
-				filtered = append(filtered, id)
+		newScores := make(map[int]int)
+		for docID, currentScore := range scores {
+			if nextCount, found := lookup[docID]; found {
+				newScores[docID] = currentScore + nextCount
 			}
 		}
+		scores = newScores
+	}
 
-		intersection = filtered
-
+	var results []Result
+	for docID, score := range scores {
+		results = append(results, Result{DocID: docID, Score: score})
 	}
-	var results []string
-	for _, id := range intersection {
-		if title, exists := idx.DocNames[id.DocId]; exists {
-			results = append(results, title)
+
+	slices.SortFunc(results, func(a, b Result) int {
+		return b.Score - a.Score
+	})
+
+	var titles []string
+	for _, res := range results {
+		if title, exists := idx.DocNames[res.DocID]; exists {
+			titles = append(titles, title)
 		}
 	}
 
-	return results
+	return titles
 }