N state probability in an M/M/2 queue I'm no ace when it comes to mathematical problems, so I've decided to go over some old college coursework and put it to code. The function works just fine and the returned values are spot on, but I'm wondering if there's a better way to solve this problem, perhaps more readable.

Problem:


  Restaurant 'A' has recently opened a 'drive through' service for its customers. The checkout has a mean arrival time of 20 per hour, with a service rate of 29 per hour. Due to the success of the new service, a second checkout has been opened.
  
  Rival restaurant, 'B', has launched an identical 'drive through' service, consisting of two checkout windows. This restaurant's mean arrival time is 19 per hour and the mean service rate is 17 per hour.
  
  In restaurant 'A', what is the probability that there will be exactly 2 customers in the queue?
  
  In restaurant 'B', what is the probability that the queue will be empty.


Formulas used

$$idle = P_0 = \frac{2\mu - \lambda}{2\mu + \lambda}$$
$$n\ state\ probability = P_n = \frac{1}{2^{n-1}}.\frac{\lambda}{\mu}^n.P_0$$

Where \$\lambda\$ is the arrival time, and \$\mu\$ is the service rate.

Code:

def n_state_probability(arrival_time, service_rate, n):
    """Return a value representing the probability an M/M/2 queue is in state n.

    To calculate the probability that the queue is empty, let n = 2.

    Arguments:
        arrival_time: mean arrival time per hour
        service_rate: mean service rate per hour
        n: amount of customers in the queue

    """
    idle = ((2 * service_rate) - arrival_time) / ((2 * service_rate) + arrival_time)
    probability = (1 / 2 ** (n - 1)) * ((arrival_time / service_rate) ** n) * idle
    if n == 2:
        return idle + probability + ((arrival_time / service_rate) * idle)
    return probability


if __name__ == '__main__':
    n_state_probability(19, 17, 2)
    # =&gt; 0.7760984526996147
    n_state_probability(20, 29, 4)
    # =&gt; 0.01377612256456733
Surely this is a trick question? The formulas only apply once the system has been running long enough that its behaviour approximates the stationary distribution. But the question explicitly says that A has "recently opened" and that B has "launched" its service. This suggests that the restaurants may not have been running long enough to approximate the stationary distribution, and so we may not be able to deduce anything about the probability of particular states.

Alternatively, the question is poorly phrased. Nonetheless, I would expect a solution to mention the assumption that the system has been running long enough for the formulas to be (approximately) valid.
Moving on, there's a second problem with the interpretation of the question. When it asks about "the queue", are we to include the customers currently being served (if any), or not?

In the docstring you write "n: amount of customers in the queue" which suggests that your interpretation is that the customers being served are in the queue.

But then you compute the answer to "what is the probability that the queue will be empty?" by adding \$P_0 + P_1 + P_2\$, which suggests that your interpretation is  that the customers being served are not in the queue.

I would expect you to make this decision one way and stick to it.
In the formulas, \$λ\$ must be the arrival rate, not (as claimed in the post) the arrival time. 

This is clear from a simple dimensional analysis — if \$λ\$ were the arrival time then it would have dimensions of time, and so could not be added to the service rate \$μ\$ which has dimensions of time\$^{-1}\$.

Hence, in the code, the variable should be named arrival_rate, not arrival_time.
The docstring says, "mean arrival time per hour", but this is meaningless.
In fact, writing "per hour" at all is superfluous since the computation doesn't depend on the units. It would work just as well if the rates were "per minute", or "per day", or whatever, so long as they were consistent.
The interface to n_state_probability has a special case: "To calculate the probability that the queue is empty, let n = 2." The trouble with this is that now you can't use the function to compute \$P_2\$.

The Zen of Python says, "Special cases aren't special enough to break the rules". So I would remove the special case and let the caller compute \$P_0 + P_1 + P_2\$ by calling the function three times and adding the results.
The computation of \$P_n\$ involves exponentiating twice:

(1 / 2 ** (n - 1)) * ((arrival_time / service_rate) ** n) * idle


A small transformation to the equation means that you would only have to exponentiate once:

2 * ((arrival_time / (2 * service_rate)) ** n) * idle

For ease of checking the code against the formulas, I would suggest using the same names.


Revised code:

# -*- coding: utf-8 -*-

def n_state_probability(arrival_rate, service_rate, n):
    """Return the probability that an M/M/2 queue is in state n, assuming
    that the system has been running long enough that it is in the
    stationary distribution.

    Arguments:
        arrival_rate: mean arrival rate
        service_rate: mean service rate
        n: number of customers in the system

    """
    λ, μ = arrival_rate, service_rate
    assert 0 &lt;= λ &lt; 2 * μ # otherwise there is no stationary distribution
    P0 = (2 * μ - λ) / (2 * μ + λ)
    if n == 0:
        return P0
    else:
        return P0 * 2 * (λ / (2 * μ)) ** n


And then the answers are:

&gt;&gt;&gt; n_state_probability(20, 29, 4)
0.01377612256456733
&gt;&gt;&gt; sum(n_state_probability(19, 17, n) for n in range(3))
0.7760984526996149