At the technical interview for my first college internship, I was asked to solve a few simple programming problems. The last one was this question:
You can move forward in increments of one step or three steps. In how many different ways can you reach $n$ steps?
I got completely nerd sniped by this question. Answering it is easy; there is a simple DP solution (below, using Python’s functools.cache to do DP for me):
from functools import cache
@cache
def steps_dp(n):
if n in [0, 1]:
return 1
elif n < 0:
return 0
else:
return steps_dp(n-1) + steps_dp(n-3)
The intuition is that you can travel $n$ steps by either traveling $n-1$ steps and then one more, or $n-3$ steps and then three more.
But the thing that nerd sniped me was trying to find a constant-time, analytical solution to this problem. In fact, I was thinking so hard about that that I completely forgot to give the DP solution and almost missed out on the internship. Over the last six years or so, I periodically remember the steps problem and spend ten minutes or so thinking about it. I spent an hour or so with ChatGPT seeing if it could solve it using generating functions (it couldn’t, at the time). But a couple of days ago, I figured it out - and it’s actually very simple!
As an intermediate step which makes the solution more clear, note that, as with many DP algorithms, you can create a slightly more memory-efficient version of the above solution by getting the answer constructively instead of recursively:
def steps_constructive(n):
if n in [0, 1]:
return 1
elif n < 0:
return 0
curr_n = 2
last_buf = [1,1,1]
while curr_n < n:
next_n = last_buf[0] + last_buf[2]
last_buf = last_buf[1:] + [next_n]
curr_n += 1
return last_buf[-1]
Here, we notice that we don’t need to store all values of the steps function, only the last three. The intuition is that you can define a buffer that you can advance by one item (from leading item $n$ to $n+1$) using only the last three items (and actually only two of those, but you need to store all three because the middle item is needed to properly advance the buffer).
Mathematically, we notice that the advancement operator we defined in this function is linear; it is represented by the matrix:
$$ A = \begin{bmatrix} 1 & 0 & 1 \\ 1 & 0 & 0 \\ 0 & 1 & 0 \end{bmatrix} $$
The initial state is $\begin{bmatrix}1 & 1 & 1\end{bmatrix}^T$ (encoding that there is one way to reach $n=0,1,2$, and by happy coincidence giving a correct value of 2 for n=3), so to advance by $k$ steps we would simply find $A^k \begin{bmatrix}1 & 1 & 1\end{bmatrix}^T$. But of course, that isn’t so simple, because exponentiating a matrix is not always a fast operation.
In this case, though, we can verify by inspection that $A$ is a nonsingular matrix, meaning that it can be diagonalized by its eigenvector decomposition. So we will find $P$, the eigenvector matrix of $A$, and $D$, the diagonal matrix of $A$’s corresponding eigenvalues. This is not always fast (in general, the eigenvector decomposition is $O(n^3)$), but it only needs to be done once, and in our case the matrix has a constant size that does not vary with the input, so we can consider this a constant-time operation. Another way to think of this is that we could manually perform the decomposition ourselves and write the algorithm as a formula using only arithmetic and no linear algebra. Of course, that would be inelegant. The full solution looks like this:
def steps_analytical(n):
if n in [0, 1]:
return 1
elif n < 0:
return 0
A = np.array([[1,0,1],[1,0,0],[0,1,0]])
evals, P = np.linalg.eig(A)
D = np.diag(evals)
P_inv = np.linalg.inv(P)
start_vec = np.array([1,1,1]).reshape((3,1))
adv_matr = P @ np.power(D, n-2) @ P_inv
end_vec = adv_matr @ start_vec
return np.round(end_vec[0], decimals=3)
Note that this actually does a very good job of not running into floating-point errors; the np.round is there to showcase this. While the DP-based solutions are $O(n)$, this analytical solution is for all purposes $O(1)$ time and memory (technically, even raising integers to the $n$ power has a time complexity of $\log n$, but it’s a tremendously fast operation).
And with that, I’m free of my years-long sporadic obsession! I hope this was a fun little application of linear algebra for you. As an exercise, you can do the Fibonacci sequence in the same way!