def dirichlet_multinomial(
        stratum_sample_tally,
        stratum_pseudocounts,
        stratum_size,
        rs):
    """
    Return a stratum tally according to the Dirichlet multinomial
    distribution, given a stratum_sample_tally, the number of votes cast
    in the stratum, and a random state.  The tally includes the
    stratum_sample_tally, so the returned tally sums to stratum_size.
    This routine is equivalent to the "Restore" procedure described
    in the "Bayesian Election Audits in One Page" note cited above.

    Input Parameters:

         stratum_sample_tally  -- a list of integers, where the i'th element
                                  is the number of votes that choice i received
                                  in the sample.
                                  (Here i indexes into actual_choices.)

         stratum_pseudocounts -- a list of integers of exactly the same length
                                 as stratum_sample_tally, giving as the i-th
                                 element the Bayesian pseudocount for the i-th
                                 choices, as a way of defining the prior.

         stratum size         -- an integer equal to the number of
                                 votes cast in this contest in this stratum.

         rs                   -- a Numpy RandomState object that is used for
                                 any random functions in the simulation of
                                 the nonsample votes. In particular,
                                 the gamma functions are made deterministic
                                 using this state.

    Returns:

        multinomial_sample       -- a list of integers, which sums up
                                    to (stratum_size - stratum_sample_size).
                                    The i'th element is equal to the
                                    simulated number of votes for choice i in
                                    the remaining, unsampled votes.

    """

    stratum_sample_size = sum(stratum_sample_tally)

    if stratum_sample_size > stratum_size:
        raise ValueError("stratum_size {} less than stratum_sample_size {}."
                         .format(stratum_size, stratum_sample_size))

    nonsample_size = stratum_size - stratum_sample_size

    stratum_sample_tally_plus_pseudocounts = \
        stratum_sample_tally + stratum_pseudocounts

    gamma_sample = [rs.gamma(k)
                    for k in stratum_sample_tally_plus_pseudocounts]
    gamma_sample_sum = float(sum(gamma_sample))
    gamma_sample = [k / gamma_sample_sum for k in gamma_sample]

    multinomial_sample_tally = rs.multinomial(nonsample_size, gamma_sample)

    # print("stratum_sample_tally:", stratum_sample_tally)
    # print("stratum_pseudocounts:", stratum_pseudocounts)
    # print(multinomial_sample_tally)
    return multinomial_sample_tally + stratum_sample_tally