def set_color_mask(A, M, c):
    """ Given image array A (h, w, 3) and mask M (h, w, 1),
        Apply color c (3,) to pixels in array A at locations M
        """
    for i in range(3):
        A_i = A[:,:,i]
        A_i[M]=c[i]
        A[:,:,i] = A_i
    return A