Problem 783: Urns

Problem 783 source code

C++

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <vector>

using ld = long double;

static ld solve_expectation(int n, int k) {
    ld m1 = static_cast<ld>(k);
    ld m2 = static_cast<ld>(k) * static_cast<ld>(k);
    ld total = 0;

    for (int t = 1; t <= n; ++t) {
        const ld m = static_cast<ld>(n - t + 2);
        const ld a = 2.0L / m;
        const ld c = (2.0L * static_cast<ld>(2 * k - 1)) /
                     (m * (static_cast<ld>(k) * m - 1.0L));

        total += c * m2 + (a - c) * m1;

        if (t == n) {
            break;
        }

        const ld r = (m - 2.0L) / m;
        const ld s = ((m - 2.0L) * (static_cast<ld>(k) * (m - 2.0L) - 1.0L)) /
                     (m * (static_cast<ld>(k) * m - 1.0L));

        const ld next_m2 = s * m2 + (2.0L * static_cast<ld>(k) * r + r - s) * m1 +
                           static_cast<ld>(k) * static_cast<ld>(k);
        const ld next_m1 = r * m1 + static_cast<ld>(k);

        m1 = next_m1;
        m2 = next_m2;
    }

    return total;
}

static ld log_binom(int n, int r) {
    if (r < 0 || r > n) {
        return -INFINITY;
    }
    return lgammal(static_cast<ld>(n) + 1.0L) -
           lgammal(static_cast<ld>(r) + 1.0L) -
           lgammal(static_cast<ld>(n - r) + 1.0L);
}

static ld hypergeom_prob(int population, int success, int draw, int picked_success) {
    if (picked_success < 0 || picked_success > success ||
        picked_success > draw || draw - picked_success > population - success) {
        return 0.0L;
    }
    const ld lp = log_binom(success, picked_success) +
                  log_binom(population - success, draw - picked_success) -
                  log_binom(population, draw);
    return expl(lp);
}

static ld brute_expectation_small(int n, int k) {
    std::vector<ld> dist(k * n + 1, 0.0L), next_dist(k * n + 1, 0.0L);
    dist[0] = 1.0L;

    ld answer = 0.0L;

    for (int t = 1; t <= n; ++t) {
        std::fill(next_dist.begin(), next_dist.end(), 0.0L);
        const int population = k * (n - t + 2);

        for (int y_prev = 0; y_prev < static_cast<int>(dist.size()); ++y_prev) {
            const ld prob_state = dist[y_prev];
            if (prob_state == 0.0L) {
                continue;
            }

            const int x = y_prev + k;
            for (int b = 0; b <= 2 * k; ++b) {
                const ld p = hypergeom_prob(population, x, 2 * k, b);
                if (p == 0.0L) {
                    continue;
                }
                answer += prob_state * p * static_cast<ld>(b) * static_cast<ld>(b);
                const int y_next = x - b;
                next_dist[y_next] += prob_state * p;
            }
        }

        dist.swap(next_dist);
    }

    return answer;
}

int main() {
    const ld sample = solve_expectation(2, 2);
    assert(fabsl(sample - 9.6L) < 1e-12L);

    const ld check_1 = solve_expectation(3, 2);
    const ld brute_1 = brute_expectation_small(3, 2);
    assert(fabsl(check_1 - brute_1) < 1e-11L);

    const ld check_2 = solve_expectation(5, 1);
    const ld brute_2 = brute_expectation_small(5, 1);
    assert(fabsl(check_2 - brute_2) < 1e-11L);

    const ld ans = solve_expectation(1'000'000, 10);
    std::cout << static_cast<long long>(llround(ans)) << '\n';
    return 0;
}

Python

import decimal

def solve_expectation(n, k):
    decimal.getcontext().prec = 50
    m1 = decimal.Decimal(k)
    m2 = decimal.Decimal(k * k)
    total = decimal.Decimal(0)
    
    dec_k = decimal.Decimal(k)
    dec_2 = decimal.Decimal(2)
    dec_1 = decimal.Decimal(1)
    
    for t in range(1, n + 1):
        m = decimal.Decimal(n - t + 2)
        a = dec_2 / m
        c = (dec_2 * decimal.Decimal(2 * k - 1)) / (m * (dec_k * m - dec_1))
        
        total += c * m2 + (a - c) * m1
        
        if t == n:
            break
            
        r = (m - dec_2) / m
        s = ((m - dec_2) * (dec_k * (m - dec_2) - dec_1)) / (m * (dec_k * m - dec_1))
        
        next_m2 = s * m2 + (dec_2 * dec_k * r + r - s) * m1 + dec_k * dec_k
        next_m1 = r * m1 + dec_k
        
        m1 = next_m1
        m2 = next_m2
        
    return total

def solve():
    ans = solve_expectation(1000000, 10)
    # Round to nearest integer (long long in C++)
    # rounding behavior: half to even is default typically, but half up is common for simple rounding
    return str(int(ans.to_integral_value(rounding=decimal.ROUND_HALF_UP)))

if __name__ == "__main__":
    print(solve())

Java

import java.math.BigDecimal;
import java.math.MathContext;
import java.math.RoundingMode;

public class Euler783 {

    public static String solveExpectation(int n, int k) {
        MathContext mc = new MathContext(50);
        BigDecimal m1 = new BigDecimal(k);
        BigDecimal m2 = new BigDecimal((long) k * k);
        BigDecimal total = BigDecimal.ZERO;

        BigDecimal decK = new BigDecimal(k);
        BigDecimal dec2 = new BigDecimal(2);
        BigDecimal dec1 = BigDecimal.ONE;
        BigDecimal dec2kMinus1 = new BigDecimal(2L * k - 1);

        for (int t = 1; t <= n; ++t) {
            BigDecimal m = new BigDecimal(n - t + 2);
            BigDecimal a = dec2.divide(m, mc);

            BigDecimal cNumerator = dec2.multiply(dec2kMinus1, mc);
            BigDecimal cDenominator = m.multiply(decK.multiply(m, mc).subtract(dec1, mc), mc);
            BigDecimal c = cNumerator.divide(cDenominator, mc);

            total = total.add(c.multiply(m2, mc), mc).add(a.subtract(c, mc).multiply(m1, mc), mc);

            if (t == n) {
                break;
            }

            BigDecimal r = m.subtract(dec2, mc).divide(m, mc);

            BigDecimal sNumerator = m.subtract(dec2, mc).multiply(
                    decK.multiply(m.subtract(dec2, mc), mc).subtract(dec1, mc), mc);
            BigDecimal sDenominator = cDenominator; // m * (k * m - 1)
            BigDecimal s = sNumerator.divide(sDenominator, mc);

            BigDecimal nextM2Term1 = s.multiply(m2, mc);
            BigDecimal nextM2Term2 = dec2.multiply(decK, mc).multiply(r, mc).add(r, mc).subtract(s, mc).multiply(m1,
                    mc);
            BigDecimal nextM2Term3 = decK.multiply(decK, mc);
            BigDecimal nextM2 = nextM2Term1.add(nextM2Term2, mc).add(nextM2Term3, mc);

            BigDecimal nextM1 = r.multiply(m1, mc).add(decK, mc);

            m1 = nextM1;
            m2 = nextM2;
        }

        // Round using half up to match typical llround behavior
        return total.setScale(0, RoundingMode.HALF_UP).toPlainString();
    }

    public static String solve() {
        return solveExpectation(1000000, 10);
    }

    public static void main(String[] args) {
        System.out.println(solve());
    }
}