The following is also not correct.

    if k not in mem:
        mem += [k]
    if k in mem:      # not the same than "else" here
        if np.random.rand() > p:
            mem.remove(k)

Your final solution is indeed correct, and I think more elegant than what we had in our paper [I am one of the authors].