matplotlib animation, assign each int a color troughout all animation

Question

I have a grid animation using matplotlib. The grid's cells (an np.array) can have 4 values throughout the animation (0, 1, 2, 3).

I want to create a legend so that during the animation: 0 = black, 1 = green, 2 = red, 3 = blue. The legend needs to stand throughout all of the animations, even if at the start some of the values are still not present in the grid.

This is my not working code:

def init(n=N, randrow=False):
    grid = np.zeros(shape=(n, n))
    grid[0] = 0
    grid[1] = 1
    grid[2] = 2
    # grid[3] = 3
    return grid


def update(d):
    global row
    row += 1
    if row >= N: return
    grid[row] = EMPTY
    grid[row + 1] = HEALTHY
    grid[row + 2] = SICK
    grid[row + 3] = RECOVERED
    mat.set_data(grid)

    return mat


if __name__ == '__main__':
    fig, ax = plt.subplots()
    ax.axis('off')
    row = 0
    grid = init()
    cmap = ListedColormap(['k', 'g', 'r', 'b'])
    mat = ax.matshow(grid, cmap=cmap)
    ani = animation.FuncAnimation(fig, update, frames=99, interval=300, save_count=50, repeat=False)
    plt.show()

The problem with the above is that since there is no 3 at the start, it maps everything using only 3 colors, and it keeps that thought-out the animation.

Does anyone know how I can fix this?