HackerRank: Construct The Array (Bottom-Up DP)

I was freshening up on my coding repertoire, and while practicing some dynamic programming (DP) questions, I came across this one on HackerRank. Usually, I adopt a top-down (memoization) approach to DP problems but in this case I made a focused effort to use a bottom-up (tabularization) approach and am documenting the same here.

Naive Attempt

Define an array mem[i][j], where $0 \leq i < n, 0 \leq j < k$ (zero-indexed), which represents the number of ways an array of length $i + 1$ can be constructed, such that it ends in $j + 1$. If constructed properly, the answer for the problem will be given by mem[n-1][x-1]. Note that, mem[0][0] = 1 and mem[0][j] = 0 $\forall j = 1, \ldots, k$, because an array of length one must end in 1.

For some $i ~ (> 0), j$, the following relation can be derived:

\[\mathrm{mem}[i][j] = \sum_{k \neq j} \mathrm{mem}[i-1][k] = \mathrm{sum}(\mathrm{mem}[i-1]) - \mathrm{mem}[i-1][j]\]

This comes from the fact that the number of ways to make an array of length $i$ ending with $j$ is equal to the sum of ways to make arrays of length $i-1$ that do not end in $j$. The code for this approach is shown below:

def countArray(n, k, x):
    # initialize array with zeros
    mem = [[0] * k for _ in range(n)]

    # only one way to make array of length 1: 
    #       it has to end with 1
    mem[0][0] = 1 

    for i in range(1,n):
        for j in range(k):
            mem[i][j] = sum(mem[i-1]) - mem[i-1][j]
    
    return mem[n-1][x-1] % ((10 ** 9) + 7)

This has memory and time complexity of $\mathcal{O}\left(nk\right)$ (if we ignore the sum operation).

First Improvement

The naive attempt instantly ran into memory errors and understadably so. A quick improvement can be made by reducing the $n\times k$ array to a $1 \times k$ array, bringing the memory complexity down to $\mathcal{O}(k)$. At the $i$-th iteration, mem[j] now will keep track of the number of ways to make an array of length $i$, which ends in $j$. This improved version can be implemented as:

def countArray_v1(n, k, x):
    # initialize array with zeros
    mem = [0] * k
    # first iteration corresponds to array of length 1,
    # only one way to make array of length 1 that ends in 1
    mem[0] = 1

    for i in range(1,n):
        mem = [sum(mem) - x for x in mem]
    return mem[x-1] % ((10 ** 9) + 7)

While the memory complexity of this is reduced to $\mathcal{O}(k)$, the time complexity is still $\mathcal{O}(nk)$ and so we still run into TLE errors.

Final Improvement

The trick to solving this problem completely can only be arrived at if you work this out on paper. For illustration, here is the output, for $k = 4$, of mem from the previous code for the first few iterations:

[1, 0, 0, 0]
[0, 1, 1, 1]
[3, 2, 2, 2]
[6, 7, 7, 7]

Clearly, $\forall j = 1, \ldots, k-1$, dp[i][j] are the same. Thus, we really only need a $1 \times 2$ array to find the answer. From the recurrence relation derived previously, we have: \(\begin{align} \mathrm{mem}[0] &= (k - 1) \cdot \mathrm{mem}[1] \\ \mathrm{mem}[1] &= (k - 2) \cdot \mathrm{mem}[1] + \mathrm{mem}[0] \end{align}\)

The code for this improved version is:

def countArray_v2(n, k, x)
    mem = [1,0]
    
    for i in range(1,n):
        mem[0], mem[1] = (k - 1) * mem[1], mem[0] + (k - 2) * mem[1]
    
    ans = mem[0] if x == 1 else mem[1]
    return ans % ((10 ** 9) + 7)

This clears all test-cases, barring one, which I have no idea how to surpass.