gistfile1.py

%matplotlib inline
import numpy as np
from scipy.spatial import Voronoi, voronoi_plot_2d
import matplotlib.pyplot as plt

points = np.array([[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2],
                    [2, 0], [2, 1], [2, 2]])

clusters = [1, 1, 1, 1, 2, 2, 2, 2, 2]

def grouped_voronoi(points, clusters):
    vor = Voronoi(points)

    # remove ridges between same groups
    removes=[]
    for i,x in reversed(list(enumerate(vor.ridge_points))):
        if (clusters[x[0]] == clusters[x[1]]):
            vor.ridge_vertices[i]=[0,0]
            #vor.ridge_vertices.pop(i)

    # remove unused vertices
    rmvert = []
    usedvertices = set([index for ridge in vor.ridge_vertices for index in ridge])
    for i in range(vor.vertices.shape[0]):
        if not i in usedvertices:
            rmvert.append(i)
    vor.vertices = np.delete(vor.vertices,rmvert,0)
    return vor

voronoi_plot_2d(grouped_voronoi(points,clusters))
plt.show()