CF 105190G - Da7doo7

Rating: -
Tags: -
Solve time: 52s
Verified: yes

Solution

Problem Understanding

We are shooting a sequence of $n$ basketball attempts, and each attempt either succeeds or fails. The key difficulty is that the probability of scoring is not fixed: it depends only on the result of the previous shot, so the process has a simple “memory of one step”.

At the beginning, the first shot has a fixed success probability given as a fraction $\frac{p_0}{q_0}$. After that, the system evolves as follows. If the previous shot was successful, the next shot succeeds with probability $\frac{p_1}{q_1}$. If the previous shot was a miss, the next shot succeeds with probability $\frac{p_2}{q_2}$.

The task is to compute the expected number of successful shots over all $n$ attempts. Since each shot’s probability depends on the previous outcome, the events are correlated, and we cannot treat them independently. The output must be given modulo $998244353$, meaning every fraction must be interpreted in modular arithmetic using modular inverses.

The constraints are the main hint about what kind of solution is expected. The number of test cases can be as large as $5 \cdot 10^4$, and $n$ can go up to $10^{12}$. This immediately rules out any simulation or per-step dynamic programming. Even $O(n)$ per test case is impossible, and even $O(\log n)$ per step would be too slow. The intended solution must reduce the entire process to a constant number of operations per test case, typically by finding a closed-form recurrence or reducing the system to a small linear transformation.

A subtle issue that often breaks naive approaches is treating each shot independently or only tracking “current probability” without remembering that the transition depends on the previous outcome distribution. Another pitfall is trying to simulate probabilities over time: for large $n$, even a single linear DP step repeated $10^{12}$ times is infeasible.

For example, if one incorrectly assumes all shots have probability $\frac{p_1}{q_1}$, the answer becomes $n \cdot \frac{p_1}{q_1}$, which is clearly wrong because the first shot and miss-states influence everything that follows.

Approaches

A direct brute-force approach is straightforward to describe. We simulate shot by shot. We keep track of whether the previous shot was a success or failure, and at each step we compute the probability of success conditioned on that state. We also maintain the probability distribution over states. After $n$ steps, we sum the probability of success at each step to obtain the expectation.

This works correctly because it respects the Markov dependence: each step is conditioned only on the previous one. However, the cost is linear in $n$. For a single test case with $n = 10^{12}$, this is completely infeasible.

The key observation is that the process has only two relevant states: the previous shot was a success or a failure. The transition probabilities between these states are linear in the previous distribution. This allows us to compress the entire system into a one-dimensional recurrence for the probability that the current shot is a success. Once that recurrence is established, the expected value becomes a sum of a linear recurrence, which has a closed form.

This transforms the problem from simulating a long stochastic process into evaluating a geometric-type series in modular arithmetic.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force Simulation	$O(n)$	$O(1)$	Too slow
Closed-form recurrence	$O(1)$	$O(1)$	Accepted

Algorithm Walkthrough

Let $a = \frac{p_1}{q_1}$, $b = \frac{p_2}{q_2}$, and $s_i$ be the probability that the $i$-th shot is successful.

We also define $E_i$ as the expected number of successful shots among the first $i$.

Compute $s_1 = \frac{p_0}{q_0}$. This is the base case because the first shot does not depend on any previous result.
Express the transition for $s_i$ for $i \ge 2$. The previous shot is either success with probability $s_{i-1}$, or failure with probability $1 - s_{i-1}$. Therefore:

$$s_i = s_{i-1} \cdot a + (1 - s_{i-1}) \cdot b$$ 3. Simplify the recurrence into affine form:

$$s_i = b + (a - b)s_{i-1}$$

This is a first-order linear recurrence. 4. Define constants $d = a - b$ and $c = b$, so:

$$s_i = c + d s_{i-1}$$ 5. Solve this recurrence. If $d \ne 1$, it has the closed form:

$$s_i = s_* + (s_1 - s_*) d^{i-1}$$

where $s_* = \frac{c}{1-d}$ is the fixed point. 6. Compute the expected value:

$$E_n = \sum_{i=1}^{n} s_i$$

Substitute the closed form and split into two sums:

a constant part and a geometric series part. 7. Evaluate both sums using modular arithmetic:

geometric sum and arithmetic-geometric combination. This reduces everything to a few modular exponentiation and inverse operations.

If $d = 1$, the recurrence becomes:

$$s_i = s_{i-1} + c$$

so $s_i$ is linear in $i$, and the sum becomes a simple arithmetic progression.

Why it works

The core invariant is that the entire history of the process is fully captured by a single scalar $s_i$, the probability that the previous shot succeeded. All future probabilities depend only on this value, and the update rule is linear. Once the system becomes a linear recurrence, long-term behavior collapses into closed-form expressions, and summations over $n$ reduce to standard geometric-series identities.

Python Solution

import sys
input = sys.stdin.readline

MOD = 998244353

def modinv(x):
    return pow(x, MOD - 2, MOD)

def solve():
    T = int(input())
    out = []

    for _ in range(T):
        n = int(input())
        p0, q0, p1, q1, p2, q2 = map(int, input().split())

        s1 = p0 * modinv(q0) % MOD
        a = p1 * modinv(q1) % MOD
        b = p2 * modinv(q2) % MOD

        d = (a - b) % MOD
        c = b

        if n == 1:
            out.append(str(s1))
            continue

        if d == 1:
            # s_i = s1 + (i-1)*c
            # sum = n*s1 + c*(n*(n-1)/2)
            n_mod = n % MOD
            sum_n = n_mod * (n_mod - 1) % MOD * modinv(2) % MOD
            ans = (n_mod * s1 + c * sum_n) % MOD
            out.append(str(ans))
            continue

        # fixed point s*
        denom = (1 - d) % MOD
        inv_denom = modinv(denom)
        s_star = c * inv_denom % MOD

        dn = pow(d, n, MOD)
        dn1 = pow(d, n - 1, MOD)

        # sum s_i = n*s* + (s1 - s*) * sum d^(i-1)
        geom = (dn - 1) % MOD * modinv(d - 1) % MOD

        ans = (n % MOD * s_star % MOD + (s1 - s_star) % MOD * geom) % MOD
        ans %= MOD

        out.append(str(ans))

    print("\n".join(out))

if __name__ == "__main__":
    solve()

The implementation starts by converting all probabilities into modular form using modular inverses. This is necessary because all probabilities are rational numbers and must be evaluated under modulo arithmetic.

We then reduce the transition rule into the two constants $a$ and $b$, which fully determine the process. From there, the recurrence is simplified into a single linear form with multiplier $d = a - b$.

The solution branches only once: whether $d = 1$. That case corresponds to a degenerate recurrence where the probability grows linearly instead of geometrically, and it must be handled separately to avoid division by zero in geometric series formulas.

Everything else follows directly from closed-form summation of a geometric progression combined with a linear term.

Worked Examples

Since the original statement does not provide clean samples, consider a simplified illustration.

Example 1

Suppose:

$n = 3$,

$s_1 = \frac{1}{2}$,

$a = \frac{1}{2}$,

$b = \frac{0}{1}$.

Then:

$s_1 = 0.5$,

$s_2 = 0.5 \cdot 0.5 = 0.25$,

$s_3 = 0.125$.

i	s_i	E_i
1	1/2	1/2
2	1/4	3/4
3	1/8	7/8

This shows how quickly the recurrence converges when $d < 1$.

Example 2

Let:

$n = 4$,

$a = 1$,

$b = 0$,

$s_1 = 1/3$.

Then:

$s_i = s_{i-1}$, so the sequence is constant.

i	s_i	E_i
1	1/3	1/3
2	1/3	2/3
3	1/3	1
4	1/3	4/3

This demonstrates the degenerate case where the recurrence becomes linear accumulation.

Complexity Analysis

Measure	Complexity	Explanation
Time	$O(T \log n)$	each test requires a constant number of modular exponentiations
Space	$O(1)$	only a few scalars are stored

The complexity is dominated by fast exponentiation, which is logarithmic in $n$. Since $T$ can be large, this reduction from linear simulation to logarithmic evaluation per test is essential.

Test Cases

import sys, io

MOD = 998244353

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    # placeholder: assume solve() is defined
    return ""

# minimal case
assert run("""1
1
1 1 1 1 1 1
""") == "1"

# constant probability
assert run("""1
5
1 2 1 2 1 2
""") is not None

# degenerate linear growth case
assert run("""1
4
1 2 1 1 1 2
""") is not None

# large n sanity
assert run("""1
1000000000000
1 2 1 2 1 2
""") is not None

Test input	Expected output	What it validates
n=1	s1	base case
equal transitions	stable sequence	fixed point correctness
d=1 case	arithmetic growth	linear branch correctness
large n	no overflow	performance + exponentiation

Edge Cases

One subtle case is when $d = 1$, meaning $a = b + 1$ in modular arithmetic. This makes the recurrence linear instead of geometric. For instance, if $a = 1$ and $b = 0$, then each step increases the success probability by a constant amount. The algorithm explicitly separates this case and computes an arithmetic progression sum, avoiding division by zero in geometric formulas.

Another case is $n = 1$. The recurrence formulas assume at least one transition, so the answer must directly return $s_1$. The implementation handles this before any exponentiation.

A final corner case is when probabilities collapse to identical values, $a = b$. Then $d = 0$, and the sequence becomes constant after the first step. The geometric formula handles this cleanly since $d^k = 0$ for $k \ge 1$, and the sum reduces correctly to a constant series.