import numpy as np

pages = np.array([
    [[1, 0], [2, 1]],
    [[5, 1], [4, 2]],
    [[0, 6], [1, 5]],
    [[-2, -1], [-1, 0]],
])
query = np.array([1.0, 0.5])

page_min = pages.min(axis=1)
page_max = pages.max(axis=1)
page_scores = np.maximum(query * page_min, query * page_max).sum(axis=1)
selected_pages = np.argsort(page_scores)[-2:]

for page, score in enumerate(page_scores):
    print(f"page {page} score: {score:.1f}")

print("selected pages:", np.sort(selected_pages).tolist())
print("KV entries read:", len(selected_pages) * 2, "of", pages.size // 2)
print("KV entries still stored:", pages.size // 2)

# Output:
"""
page 0 score: 2.5
page 1 score: 6.0
page 2 score: 4.0
page 3 score: -1.0
selected pages: [1, 2]
KV entries read: 4 of 8
KV entries still stored: 8
"""