 Solution

One simulation contributes its largest observed infected population and its last recorded time. Repeating this operation produces one row per realization:

def sir_sim_stats(n_sim,t_max,X0,beta,gamma,rng=None):
    result=np.empty((n_sim,2))
    for run in range(n_sim):
        t,s=sir_sim(t_max,X0,beta,gamma,rng)
        result[run]=s[:,1].max(),t[-1]
    return result

The maximum must be taken over the infected column, states[:,1], rather than across all compartments. With the implementation in part (a), the final time is the epidemic extinction time if extinction occurs first, and otherwise the specified horizon; an outbreak still active at that horizon has a censored extinction time.