Conway's Game of Life: How to colour the cells based on number of neighbours? (change number of neighbours to age of cells)

The number of neighbours is already calculated in the animate function for the last generation as "newu". I have introduced an "age map", which gets capped at 4. My insertions are in the #addon comments.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import ListedColormap
def createuniverse(N=50, p=0.5):
    return np.random.choice([0, 1], size=(N, N), p=[1-p, p])
N = 50
universe = createuniverse(N=N, p=0.5)
generations = 0
addon: inital age & target colors 0 - 4

age = np.zeros(universe.shape)
colors = ['white', 'green', 'yellow', 'orange', 'red']
discretecmap = ListedColormap(colors)
def animate(frame, universe, img):
    # addon: added age as global variable for modification
    global generations, age
    newu = np.zeros((N, N), dtype = int)
    universe[0,:] = universe[-1,:] = universe[:,0] = universe[:,-1] = 0
    newu[1:-1,1:-1] = (universe[ :-2, :-2] + universe[ :-2,1:-1] + universe[ :-2,2:] +
                        universe[1:-1, :-2]                + universe[1:-1,2:] +
                        universe[2:  , :-2] + universe[2:  ,1:-1] + universe[2:  ,2:])
    birth = (newu==3)[1:-1,1:-1] & (universe[1:-1,1:-1]==0)
    survive = ((newu==2) | (newu==3))[1:-1,1:-1] & (universe[1:-1,1:-1]==1)
    universe[:] = np.zeros((N, N), dtype = int)
    universe[1:-1,1:-1][birth | survive] = 1
    population = np.countnonzero(universe)
    # addon: calculate & plot age of cells
    age += universe
    age *= universe # filter dead cells
    age[age>4] = 4 # cap the age at 4
    img.setdata(age)
    ax.setxlabel('generations = {}, population = {}'.format(generations, population))
    generations += 1
fig = plt.figure(figsize=(5, 5))
ax = plt.axes()
ax.setyticklabels([])
ax.setxticklabels([])
img = ax.imshow(universe, vmin=0, vmax=4, cmap=discrete_cmap)
ani = FuncAnimation(fig, animate, fargs=(universe, img,),
                    frames=200, interval=200)
plt.show()