Histogram different colours for clusters

Question

So I am currently trying to learn matplotlib using the iris dataset and its going pretty well however when making a matrix of scatter plots with histograms going down the diagonal I am trying to get the histograms to also change colours depending on clusters.

So if anyone knows how to do this easily it would be appreciated.

I had one person say I could pass in 3 subsets? Something like this maybe?

ax[feature1,feature2].hist(x=(subset1,subset2,subset3))

However I have not had any success with trying to get it to work. I hope I'm not being stupid as I've not used matplotlib before.

Here is my current code which makes a matrix of scatter plots perfectly fine.

clusterModel = KMeans(n_clusters=K)
clusterModel.fit(irisX)
cluster_ids = clusterModel.predict(irisX)


def show_scatterplot_matrix(X,y,featureNames,title=None):
    f = X.shape[1]
    if(len(y) != X.shape[0]):
        print("Error, the y array must have the same length as there are rows in X")
        return
    fig, ax = plt.subplots(f,f,figsize=(12,12))
    plt.set_cmap('jet')
    for feature1 in range(f):
        ax[feature1,0].set_ylabel( featureNames[feature1])
        ax[0,feature1].set_xlabel( featureNames[feature1])
        ax[0,feature1].xaxis.set_label_position('top') 
        for feature2 in range(f):
            xdata = X[:,feature1]
            ydata = X[:,feature2]
            if (feature1 != feature2):
                ax[feature1, feature2].scatter(xdata,ydata,c=y)
            else:
                ax[feature1, feature2].hist(xdata)

    if title != None:
        fig.suptitle(title,fontsize=16,y=0.925)

show_scatterplot_matrix(irisX,cluster_ids,feature_names,title="Visualisation of the Iris Data")

Here is also the current output of the data as a image: Matrix of scatter plots