from numpy import random, array, float
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from sklearn.preprocessing import scale


# create fake income/age clusters for N people in k clusters
def create_clustered_data(N, k):
    random.seed(10)
    points_per_cluster = float(N) / k
    x = []
    for i in range(k):
        income_centroid = random.uniform(20000.0, 200000.0)
        age_centroid = random.uniform(20.0, 70.0)
        for j in range(int(points_per_cluster)):
            x.append([random.normal(income_centroid, 10000.0), random.normal(age_centroid, 2.0)])
    x = array(x)
    return x


data = create_clustered_data(100, 5)
model = KMeans(n_clusters=5)

model = model.fit(scale(data))
print(model.labels_)

plt.figure(figsize=(8, 6))
plt.scatter(data[:, 0], data[:, 1], c=model.labels_.astype(float))
plt.show()