公式

F - Random Vault Heist 解説 by en_translator


This problem asks to repeatedly choose randomly from the remaining until a condition is violated. This kind of problem may be understood well by rearrange randomly, and pick one by one from the beginning until the condition is violated.


Let \(T\) be the set of safes, and \(A_{\text{all}} = \sum_{i \in T} A_i \).

Rephrase the condition for the robber to continue stealing as:

Let \(S\) be the set of safes stolen so far. If \(\sum_{i \in S} A_i \lt X\), he opens the next safe.

Then the sought expected value can be reinterpreted as follows:

The sum of

  • the probability of reaching the state where “set \(S\) has been stolen” \(\times\) the expected amount of money stolen next

over all \(S\) such that \(\sum_{i \in S} A_i \lt X\).

The former is the probability that, when the \(N\) safes are arranged in a line, specific \(|S|\) items form a prefix. This is \(\displaystyle \frac{1}{\binom{N}{|S|}}\).

The latter is chosen uniformly at random from the safes other than \(S\), so \(\displaystyle \frac{A_{\text{all}} - \sum_{i \in S} A_i}{N - |S|}\).


We want to evaluate this value for all \(S\) with \(\sum_{i \in S} A_i \lt X\), but there are at most \(2^N\) ways to choose \(S\), so enumerating them all within the execution time limit is impossible.

However, for sets with the same \(|S| = k\), the denominator is common. Therefore, it suffices to compute for all \(k\):

  • the number of sets of size \(k\), and
  • the sum of \(\sum_{i \in S} A_i\) for all such sets.

Finding this directly is still difficult, but the following technique helps us evaluating it fast.


We use the meet-in-the-middle trick. Divide \(T\) into two sets, \(T_l\) and \(T_r\).

We will evaluate the sought values for all \(S\) that chooses a set \(S_l\) from \(T_l\), and \(k\) elements from \(T_r\). Then it is sufficient if we can, given any \(S_r \subset T_r\) such that

  • \(|S_r| = k\),
  • \(\sum_{i \in S_r} A_i \lt X - \sum_{i \in S_l} A_i\),

obtain

  • the number of such sets \(S_r\), and
  • the sum of \(\sum_{i \in S_r} A_i\) for all such sets.

This can be obtained by precalculating, for each \(|S_r| = k\):

  • the sequence of \(\sum_{i \in S_r} A_i\) arranged in ascending order, and
  • its prefix sum.

by binary searching. All that left is to find the expected value for all \((S_l, k)\) and take the sum.


Let us consider the complexity. Let \(l = |T_l|\) and \(r = |T_r|\).

The precalculation requires enumerating \(2^r\) elements and sorting for set size \(k\), costing \(O(r \cdot 2 ^ {r})\). For each of the \(2^l\) sets \(S_i\) and \(k ~ (0 \le k \le r)\), we need to binary search, costing \(O(2^l \cdot r^2)\) time.

By taking both \(l\) and \(r\) as \(N/2\), the problem can be solved in a total of \(O(N^2 2^{N/2})\) time.


By monotonicity, the binary search can be turned into sliding-window scanning, allowing \(O(N 2^{N/2})\) solution.

If you find the time limit tight, taking larger \(r\) might help, as the expression \(O(r2^r + r^22^l)\) suggests.


Sample code (PyPy)

from math import comb
from bisect import bisect_left

n, x = map(int, input().split())
a = list(map(int, input().split()))

mod = 998244353

def gen(a):
    n = len(a)
    r = [[] for _ in range(n + 1)]
    r[0].append(0)
    for x in a:
        for k in reversed(range(n)):
            if r[k]: r[k+1].extend(s + x for s in r[k])
    return r

h = n // 2
tls, trs = gen(a[:h]), gen(a[h:])

a_all = sum(a) % mod
q = [pow((n - k) * comb(n, k), -1, mod) for k in range(n)] + [0] # denominator value

ans = 0
for k, tr in enumerate(trs):
    tr.sort()
    
    pre = [0] * (len(tr) + 1) # prefix sum
    for i in range(len(tr)):
        pre[i+1] = (pre[i] + tr[i]) % mod
    
    for i, tl in enumerate(tls):
        for sl in tl:
            c = bisect_left(tr, x - sl)
            p = (c * (a_all - sl % mod) - pre[c]) % mod
            ans += q[i+k] * p % mod
            ans %= mod

print(ans)

投稿日時:
最終更新: