Problem: balance_scale (2022 round2)

Your friend baked \(N\) batches of cookies, the first with chocolate chips, and the rest with raisins. Batch \(i\) has \(C_i\) cookies all having the same weight \(W_i\). Cookies in different batches may have the same weight. Initially, all cookies are scattered on the table, and you hope to find a hefty one to eat with your trusty balance scale.

First, you select a single cookie uniformly at random from the table and place it on the left side of the scale. Then, \(K\) times, you will select a cookie uniformly at random from the table, place it on the empty side of the scale, and throw away the lighter of the two cookies on the scale (or the cookie on the left side if they both weigh the same).

{{PHOTO_ID:502174485028685|WIDTH:700}}

Determine the probability that after throwing away \(K\) cookies, the remaining cookie on the scale is a yummy chocolate chip cookie from batch \(1\).

If we write this probability as a quotient of integers \(p/q\) in lowest terms, then you should output this quotient modulo \(1{,}000{,}000{,}007\) — in other words, output the unique integer \(x\) such that \(0 \le x \lt 1{,}000{,}000{,}007\) and \(p = (x * q) \text{ mod } 1{,}000{,}000{,}007\).


# Constraints

\(1 \le T \le 1{,}000\)
\(2 \le N \le 3{,}000\)
\(1 \le K \lt \sum_{i=1}^{N} C_i\)
\(1 \le C_i \le 3{,}000\)
\(1 \le W_i \le 10^{9}\)


# Input Format

Input begins with a single integer \(T\), the number of test cases. For each test case, there is first a line containing two-space separated integers \(N\) and \(K\). Then, \(N\) lines follow, the \(i\)th of which contains two space-separated integers \(C_i\) and \(W_i\).


# Output Format

For the \(i\)th test case, output a single line containing `"Case #i: "` followed by the unique integer \(x\) such that \(p = (x * q) \text{ mod } 1{,}000{,}000{,}007\).


# Sample Explanation

In the first case, there are \(5\) batches with \(1\) cookie each. Since we only weigh \(K=1\) time, the only way for the last cookie to be chocolate chip is if it's one of the two placed on the scale. There's a \(\frac{1}{5}\) chance of picking it for the initial cookie, and a \(\frac{4}{5}\cdot\frac{1}{4}\) chance of picking it for the second cookie. In either case, because it's the heaviest, it will remain on the scale. Adding these independent events' probabilities, we get a \(\frac{1}{5} + \frac{4}{5}\cdot\frac{1}{4} = \frac{2}{5}\) chance that the last cookie on the scale is chocolate chip. Here, we print \(800{,}000{,}006\) since \(2 = (800{,}000{,}006 * 5) \text{ mod } 1{,}000{,}000{,}007\).

The second case is the same, except we must weigh a second time. There's a \(\frac{2}{5}\) chance of picking chocolate chip the first or second times, and a \(\frac{3}{5}\cdot\frac{1}{3}\) chance of picking it the third time. The answer is \(\frac{2}{5} + \frac{3}{5}\cdot\frac{1}{3} = \frac{3}{5}\). We print \(200{,}000{,}002\) since \(3 = (200{,}000{,}002 * 5) \text{ mod } 1{,}000{,}000{,}007\).

In the third case, there are \(2\) batches of \(10\) cookies each, where batch \(2\) cookies are heavier. Since we must weigh and throw out \(10\) cookies, we must have weighed a cookie from batch \(2\) at least once, after which a batch \(1\) cookie can never remain on the scale. Therefore the answer is \(0\).

In the fourth case, there is a \(\frac{5}{24}\) chance that the final cookie on the scale is chocolate chip.


Sample input:
6
5 1
1 3000
1 2000
1 1000
1 2000
1 1000
5 2
1 3000
1 2000
1 1000
1 2000
1 1000
2 10
10 1
10 2
5 2
2 50
1 40
1 50
1 60
3 50
4 2993
3000 999999999
2995 1000000000
1552 888888888
1336 999999999
3 1
1 10
2 9
1 11


Sample output:
Case #1: 800000006
Case #2: 200000002
Case #3: 0
Case #4: 208333335
Case #5: 590307096
Case #6: 333333336


Solution: The problem is asking: "Out of a uniformly randomly chosen subset of \(K+1\) cookies, what is the probability that the heaviest cookie (with ties broken uniformly randomly) is from batch \(1\)?" Note that it's irrelevant how ties are broken by choosing the left cookie, because for any random subset lined up in a random order to be place on the scale, it's equally likely for any two cookies of the same weight to be placed on either side.

Let \(X\), \(Y\), and \(Z\) be the numbers of cookies outside of batch \(1\) that weigh less than, equal to, and greater than \(W_1\), respectively. Let \(M = C_1 + X + Y + Z\) be the total number of cookies. Note that if \(C_1 + X + Y < K + 1\), then any size \(K+1\) subset must contain at least \(1\) cookie heavier than \(W_1\), in which case the answer is \(0\). Otherwise, the answer is determined as follows:

1. There are \(\binom{C_1 + X + Y}{K+1}\) ways to pick \(K+1\) cookies that are not heavier than \(W_1\).
2. There are \(\binom{X}{K+1}\) ways to pick \(K+1\) cookies that are strictly lighter than \(W_1\).
3. The number of ways to pick at least \(1\) cookie of weight equal to \(W_1\) (some, none, or all from batch \(1\)) is the first number minus the second number.
4. The probability of picking at least \(1\) cookie of weight equal to \(W_1\) is equal to the previous number divided by \(\binom{M}{K + 1}\) (the total number of ways to pick \(K+1\) cookies). 
5. The conditional probability of a tie being broken with a batch \(1\) cookie out of all equal weight cookies picked (*given* that we picked at least \(1\) equal weight cookie) is equal to \(\frac{C_1}{C_1 + Y}\). Note that this is independent of any attributes of a specific chosen subset, as each cookie of weight equal to \(W_1\) has the same \(\frac{C_1}{C_1 + Y}\) chance of being a batch \(1\) cookie.
6. The final answer is the product of the previous two values (the conditional probability times the prior probability).

If we precompute factorials up to \(M\) in \(\mathcal{O}(M)\) time, then each binomial coefficient can be computed in time \(\mathcal{O}(1)\), plus the time to compute a mod-inverse. After that, there are only a constant number of modular arithmetic steps, which we can perform in \(\mathcal{O}(1)\) as well.

[See David Harmeyer's solution video here.](https://www.youtube.com/watch?v=rg0Q73SMhUM)


Code: #include <algorithm>
#include <cassert>
#include <iostream>
#include <vector>
using namespace std;

const int MOD = 1000000007;
using LL = long long;

class mint {
  LL val;

  mint& normalize(LL x) {
    val = x % MOD;
    if (val < 0) {
      val += MOD;
    }
    return *this;
  }

  static LL inverse(LL a) {
    LL u = 0, v = 1, m = MOD;
    while (a != 0) {
      LL t = m / a;
      m -= t * a;
      swap(a, m);
      u -= t * v;
      swap(u, v);
    }
    assert(m == 1);
    return u;
  }

 public:
  mint(LL x = 0) { normalize(x); }

  LL operator()() const { return val; }

  template <typename U>
  explicit operator U() const { return static_cast<U>(val); }

  mint& operator=(LL x) { return normalize(x); }
  mint& operator=(const mint& x) { val = x.val; return *this; }

  mint& operator+=(const mint& x) { return normalize(val + x.val); }
  mint& operator-=(const mint& x) { return normalize(val - x.val); }
  mint& operator*=(const mint& x) { return normalize(val * x.val); }
  mint& operator/=(const mint& x) { return *this *= mint(inverse(x.val)); }

  mint& operator+=(LL x) { return *this += mint(x); }
  mint& operator-=(LL x) { return *this -= mint(x); }
  mint& operator*=(LL x) { return *this *= mint(x); }
  mint& operator/=(LL x) { return *this /= mint(x); }

  mint& operator++() { return *this += 1; }
  mint& operator--() { return *this -= 1; }
  mint operator++(int) { mint z(*this); ++*this; return z; }
  mint operator--(int) { mint z(*this); --*this; return z; }

  friend mint operator+(mint x, const mint& y) { return x += y; }
  friend mint operator*(mint x, const mint& y) { return x *= y; }
  friend mint operator-(mint x, const mint& y) { return x -= y; }
  friend mint operator/(mint x, const mint& y) { return x /= y; }

  friend mint operator+(mint x, LL y) { return x += y; }
  friend mint operator*(mint x, LL y) { return x *= y; }
  friend mint operator-(mint x, LL y) { return x -= y; }
  friend mint operator/(mint x, LL y) { return x /= y; }

  friend mint operator+(LL x, mint y) { return y += x; }
  friend mint operator*(LL x, mint y) { return y *= x; }
  friend mint operator-(LL x, const mint& y) { mint z(x); return z -= y; }
  friend mint operator/(LL x, const mint& y) { mint z(x); return z /= y; }

  bool operator <(const mint& x) const { return val < x.val; }
  bool operator==(const mint& x) const { return val == x.val; }
  bool operator >(const mint& x) const { return val > x.val; }
  bool operator!=(const mint& x) const { return val != x.val; }
  bool operator<=(const mint& x) const { return val <= x.val; }
  bool operator>=(const mint& x) const { return val >= x.val; }

  bool operator <(LL x) const { return val < x; }
  bool operator==(LL x) const { return val == x; }
  bool operator >(LL x) const { return val > x; }
  bool operator!=(LL x) const { return val != x; }
  bool operator<=(LL x) const { return val <= x; }
  bool operator>=(LL x) const { return val >= x; }
};

class factorial : vector<mint> {
  void lazy_eval(int n) {
    for (LL p = size(); n >= p; ++p) {
      push_back(back() * p);
    }
  }

 public:
  factorial() { push_back(1); }

  mint choose(LL n, LL k) {
    if (n < 0 || k < 0 || n < k) {
      return 0;
    }
    if (k == 0 || k == n) {
      return 1;
    }
    if (n >= MOD && n <= k + MOD - 1) {
      return 0;
    }
    lazy_eval(n);
    return at(n) / (at(k) * at(n - k));
  }

  mint operator[](LL n) {
    lazy_eval(n);
    return at(n);
  }
};

factorial F;

LL solve() {
  int N, K, c0, w0;
  cin >> N >> K >> c0 >> w0;
  LL lighter = 0, equal = 0, heavier = 0;
  for (int i = 1, c, w; i < N; i++) {
    cin >> c >> w;
    if (w < w0) {
      lighter += c;
    } else if (w == w0) {
      equal += c;
    } else {
      heavier += c;
    }
  }
  LL total = c0 + lighter + equal + heavier;
  LL not_heavier = total - heavier;
  K++;
  mint ways_some_equal = F.choose(not_heavier, K) - F.choose(lighter, K);
  mint prob_some_equal = ways_some_equal / F.choose(total, K);
  return LL(prob_some_equal * c0 / (c0 + equal));
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  int T;
  cin >> T;
  for (int t = 1; t <= T; t++) {
    cout << "Case #" << t << ": " << solve() << endl;
  }
  return 0;
}