Search the Beach

A ring was lost on a beach, which is divided into n squares. prior[i] is the probability that the ring is in square i (they add up to 1: the ring is certainly on the beach). Searching square i finds the ring with probability detect[i] if it is there (soft sand hides it better than wet sand), and never finds it if it is not.

The search team always searches next the square where this search is most likely to succeed, given that every search so far has failed. When several squares give exactly the same chance, they take the smallest index. After each failed search, they update their beliefs about where the ring is with Bayes' rule.

Write search_plan(prior, detect, k) that returns a tuple (plan, found), assuming the first k searches all fail:

• plan is the list of the k squares searched, in order;
• found is the probability, computed before any search, that the ring is found within these k searches.

The setup provides beach(n, seed), which returns random lists (prior, detect) for n squares.

Examples

Input:  prior = [0.5, 0.3, 0.2], detect = [0.4, 0.9, 0.5], k = 3
Output: ([1, 0, 0], 0.59)
Explanation: the chances of success on the first search are 0.2, 0.27 and 0.1, so
square 1 goes first. It fails, which makes square 1 unlikely: the probabilities become
0.5/0.73, 0.03/0.73, 0.2/0.73, and square 0 is now best, twice in a row. The ring is
found on the first search with probability 0.27, on the second with 0.2 and on the
third with 0.12: 0.59 in total.

Constraints

• 1 <= n <= 10**4, 0 <= k <= 10**5
• every prior[i] is positive and they add up to 1; 0 < detect[i] <= 1
• floats are compared with a tolerance of 1e-6; the largest tests need about O(k log n)

Goals

• Update a probability map with Bayes' rule after an unsuccessful search
• Choose every search to maximise the chance of success on that search
• Notice that normalising does not change which cell is best, and use a heap