First I would choose a different representation of the colors during
the computation, to avoid the “expensive” COLORS-{a, b}).pop()
operation, which is executed
$$
(n-1) + (n-2) + \ldots + 1 = \frac{(n-1)n}{2}
$$
times for an input string of length \$ n \$. It becomes simpler if we
represent the colors as numbers instead, for example
$$
\text{Red} = 0, \text{Green} = 1, \text{Blue} = 2 \, .
$$
Now the combination of two colors a and b can be computed as
a if a == b else 3 - a - b
without requiring any set or dictionary lookup. This leads to the
following implementation:
def color_to_num(col):
""" Translate the colors 'R', 'G', 'B' to the numbers 0, 1, 2. """
return 0 if col == 'R' else 1 if col == 'G' else 2
def num_to_color(num):
""" Translate the numbers 0, 1, 2 to the colors 'R', 'G', 'B'. """
return 'R' if num == 0 else 'G' if num == 1 else 'B'
def triangle(row):
numbers = [color_to_num(col) for col in row]
while len(numbers) > 1:
numbers = [ a if a == b else 3 - a - b for a, b in zip(numbers, numbers[1:])]
return num_to_color(numbers[0])
In my test (on a 1.2 GHz Intel Core m5 MacBook) this reduces the time for
an input string with 10,000 characters from (approximately) 18 seconds to
5.5 seconds.
But one can do better. The above solution still has a \$ O(n^2) \$ complexity.
In order to get rid of that, we need some mathematics (compare
Three Color Triangle Challenge
on Mathematics Stack Exchange).
The are two key observations:
- If all (numeric) colors are viewed “modulo 3” then the combination of
\$ a \$ and \$ b \$ is just \$ -(a+b) \$.
- The final color can be computed “directly” from the initial row, using
binomial coefficients (the numbers from Pascal's triangle).
Here is an example for the first observation: The combination of Red and Blue
is computed as
$$
- (0 + 2) \bmod 3 = -2 \bmod 3 = 1
$$
and that is Green.
And an example for the second observation with an initial row of 4 colors:
$$
a_0, a_1, a_2, a_3 \\
-(a_0 + a_1), -(a_1 + a_2), -(a_2 + a_3) \\
a_0 + 2a_1 + a_2, a_1 + 2a_2 + a_3 \\
-(a_0 + 3a_1 + 3 a_2 + a_3)
$$
The general formula for an initial row of \$ n+1 \$ colors
$$
a_0, a_1, \ldots, a_n
$$
is
$$
\text{FinalColor} = (-1)^n \sum_{k=0}^n \binom{n}{k} a_k \bmod 3 \, .
$$
There is still one problem: the calculation of the binomial coefficients.
Using the definition
$$
\binom{n}{k} = \frac{n(n-1) \cdots (n-k+1)}{k!}
$$
can lead to large intermediate results and has complexity \$ O(k) \$,
so that the total complexity is still \$ O(n^2) \$.
Here we need mathematics again: Lucas's theorem
provides a way to calculate \$ \binom{n}{k} \bmod p \$ efficiently
if \$ p \$ is a prime number. Luckily, 3 is a prime number!
This gives the following code:
def binomial_mod3(n, k):
""" Compute the binomial coefficient C(n, k) modulo 3.
It is assumed that 0 <= k <= n.
The implementation uses Lucas's theorem, see for example
https://en.wikipedia.org/wiki/Lucas%27s_theorem .
"""
result = 1
while n > 0:
n3 = n % 3
k3 = k % 3
# 'Ad hoc' computation of C(n3, k3):
if k3 > n3:
return 0 # Return immediately if a factor is zero.
temp = 1 if k3 == 0 or k3 == n3 else 2
result = (result * temp) % 3
n = n // 3
k = k // 3
return result
def color_to_num(col):
""" Translate the colors 'R', 'G', 'B' to the numbers 0, 1, 2. """
return 0 if col == 'R' else 1 if col == 'G' else 2
def num_to_color(num):
""" Translate the numbers 0, 1, 2 to the colors 'R', 'G', 'B'. """
return 'R' if num == 0 else 'G' if num == 1 else 'B'
def triangle(row):
# Compute the sum of C(n, k) * col[k] modulo 3:
n = len(row) - 1
result = 0
for k, col in enumerate(row):
temp = binomial_mod3(n, k) * color_to_num(col)
result = (result + temp) % 3
# Invert the result if n is odd:
if n % 2 == 1:
result = (- result) % 3
return num_to_color(result)
Note that we iterate over the given input string, so that
no additional lists are created.
In my test this computes the final color for an input string with 10,000
characters in approximately 0.1 seconds, and for 100,000 characters in
about 0.2 seconds.
