Revert cast, improved query function

larry-he · larry-he · commit 65636ca3ecd2 · 2016-03-24T12:50:38.000-04:00
Cast was not needed since the return value was already being cast. Query
function changed to use heapq.nlargest() for simplicity
diff --git a/grid.py b/grid.py
@@ -44,4 +44,4 @@ def vector_clip(vector, lowest, highest):
     """Return vector, except if any element is less than the corresponding
     value of lowest or more than the corresponding value of highest, clip to
     those values."""
-    return type(vector)(list(map(clip, vector, lowest, highest)))
+    return type(vector)(map(clip, vector, lowest, highest))
diff --git a/text.py b/text.py
@@ -147,15 +147,17 @@ def query(self, query_text, n=10):
         qwords = [w for w in words(query_text) if w not in self.stopwords]
         shortest = argmin(qwords, lambda w: len(self.index[w]))
         docs = self.index[shortest]
-        results = sorted([(sum([self.score(w, d) for w in qwords]), d) for d in docs])
-        results.reverse()
-        return results[:n]
+        return heapq.nlargest(n, ((total_score(qwords, doc), doc) for doc in docs))
 
-    def score(self, word, docid):
+    def score(self, word, doc):
         "Compute a score for this word on this docid."
         # There are many options; here we take a very simple approach
-        return (math.log(1 + self.index[word][docid]) /
-                math.log(1 + self.documents[docid].nwords))
+        return (math.log(1 + self.index[word][doc]) /
+                math.log(1 + self.documents[doc].nwords))
+
+    def total_score(qwords, doc):
+        "Compute the sum of the scores of the queried words on this doc."
+        return sum(self.score(qword, doc) for qword in qwords)
 
     def present(self, results):
         "Present the results as a list."
