diff --git a/src/nvmath/Fitting.cpp b/src/nvmath/Fitting.cpp
index b01b18a..53bd985 100644
--- a/src/nvmath/Fitting.cpp
+++ b/src/nvmath/Fitting.cpp
@@ -3,6 +3,7 @@
 #include "Fitting.h"
 
 #include <nvcore/Algorithms.h> // max
+#include <nvcore/Containers.h> // swap
 #include <float.h> // FLT_MAX
 
 using namespace nv;
@@ -78,7 +79,7 @@ Vector3 nv::ComputePrincipalComponent(int n, const Vector3 * points, const float
 
 
 
-void nv::Compute4Means(int n, const Vector3 * points, const float * weights, Vector3::Arg metric, Vector3 * cluster)
+int nv::Compute4Means(int n, const Vector3 * points, const float * weights, Vector3::Arg metric, Vector3 * cluster)
 {
 	Vector3 centroid = ComputeCentroid(n, points, weights, metric);
 	
@@ -90,7 +91,7 @@ void nv::Compute4Means(int n, const Vector3 * points, const float * weights, Vec
 	mini = maxi = 0;
 	
 	float mindps, maxdps;
-	mindps = maxdps = dot(points[0], principal);
+	mindps = maxdps = dot(points[0] - centroid, principal);
 	
 	for (int i = 1; i < n; ++i)
 	{
@@ -106,13 +107,15 @@ void nv::Compute4Means(int n, const Vector3 * points, const float * weights, Vec
 		}
 	}
 
+	//cluster[0] = points[mini];
+	//cluster[1] = points[maxi];
 	cluster[0] = centroid + mindps * principal;
-	cluster[3] = centroid + maxdps * principal;
-	cluster[1] = (2 * cluster[0] + cluster[1]) / 3;
-	cluster[2] = (2 * cluster[1] + cluster[0]) / 3;
+	cluster[1] = centroid + maxdps * principal;
+	cluster[2] = (2 * cluster[0] + cluster[1]) / 3;
+	cluster[3] = (2 * cluster[1] + cluster[0]) / 3;
 
 	// Now we have to iteratively refine the clusters.
-	while(true)
+	while (true)
 	{
 		Vector3 newCluster[4] = { Vector3(zero), Vector3(zero), Vector3(zero), Vector3(zero) };
 		float total[4] = {0, 0, 0, 0};
@@ -141,15 +144,119 @@ void nv::Compute4Means(int n, const Vector3 * points, const float * weights, Vec
 			newCluster[j] /= total[j];
 		}
 
-		if (equal(cluster[0], newCluster[0]) && equal(cluster[3], newCluster[3]))
+		if ((equal(cluster[0], newCluster[0]) || total[0] == 0) && 
+			(equal(cluster[1], newCluster[1]) || total[1] == 0) && 
+			(equal(cluster[2], newCluster[2]) || total[2] == 0) && 
+			(equal(cluster[3], newCluster[3]) || total[3] == 0))
 		{
-			break;
+			return (total[0] != 0) + (total[1] != 0) + (total[2] != 0) + (total[3] != 0);
 		}
-		
-		// @@ We should choose the optimal assignment.
+
 		cluster[0] = newCluster[0];
+		cluster[1] = newCluster[1];
+		cluster[2] = newCluster[2];
 		cluster[3] = newCluster[3];
-		cluster[1] = (2 * cluster[0] + cluster[1]) / 3;
-		cluster[2] = (2 * cluster[1] + cluster[0]) / 3;
+
+		// Sort clusters by weight.
+		for (int i = 0; i < 4; i++)
+		{
+			for (int j = i; j > 0 && total[j] > total[j - 1]; j--)
+			{
+				swap( total[j], total[j - 1] );
+				swap( cluster[j], cluster[j - 1] );
+			}
+		}
 	}
 }
+
+/*
+int nv::Compute2Means(int n, const Vector3 * points, const float * weights, Vector3::Arg metric, Vector3 * cluster)
+{
+	Vector3 centroid = ComputeCentroid(n, points, weights, metric);
+	
+	// Compute principal component.
+	Vector3 principal = ComputePrincipalComponent(n, points, weights, metric);
+	
+	// Pick initial solution.
+	int mini, maxi;
+	mini = maxi = 0;
+	
+	float mindps, maxdps;
+	mindps = maxdps = dot(points[0] - centroid, principal);
+	
+	for (int i = 1; i < n; ++i)
+	{
+		float dps = dot(points[i] - centroid, principal);
+		
+		if (dps < mindps) {
+			mindps = dps;
+			mini = i;
+		}
+		else {
+			maxdps = dps;
+			maxi = i;
+		}
+	}
+
+	cluster[0] = points[mini];
+	cluster[3] = points[maxi];
+	//cluster[0] = centroid + mindps * principal;
+	//cluster[1] = centroid + maxdps * principal;
+	cluster[2] = (2 * cluster[0] + cluster[1]) / 3;
+	cluster[3] = (2 * cluster[1] + cluster[0]) / 3;
+
+	// Now we have to iteratively refine the clusters.
+	while (true)
+	{
+		Vector3 newCluster[4] = { Vector3(zero), Vector3(zero), Vector3(zero), Vector3(zero) };
+		float total[4] = {0, 0, 0, 0};
+		
+		for (int i = 0; i < n; ++i)
+		{
+			// Find nearest cluster.
+			int nearest = 0;
+			float mindist = FLT_MAX;
+			for (int j = 0; j < 4; j++)
+			{
+				float dist = length_squared(cluster[j] - points[i]);
+				if (dist < mindist)
+				{
+					mindist = dist;
+					nearest = j;
+				}
+			}
+			
+			newCluster[nearest] += weights[i] * points[i];
+			total[nearest] += weights[i];
+		}
+
+		for (int j = 0; j < 4; j++)
+		{
+			newCluster[j] /= total[j];
+		}
+
+		if ((equal(cluster[0], newCluster[0]) || total[0] == 0) && 
+			(equal(cluster[1], newCluster[1]) || total[1] == 0) && 
+			(equal(cluster[2], newCluster[2]) || total[2] == 0) && 
+			(equal(cluster[3], newCluster[3]) || total[3] == 0))
+		{
+			return (total[0] != 0) + (total[1] != 0) + (total[2] != 0) + (total[3] != 0);
+		}
+
+		cluster[0] = newCluster[0];
+		cluster[1] = newCluster[1];
+		cluster[2] = newCluster[2];
+		cluster[3] = newCluster[3];
+
+		// Sort clusters by weight.
+		for (int i = 0; i < 4; i++)
+		{
+			for (int j = i; j > 0 && total[j] > total[j - 1]; j--)
+			{
+				swap( total[j], total[j - 1] );
+				swap( cluster[j], cluster[j - 1] );
+			}
+		}
+	}
+}
+*/
\ No newline at end of file
diff --git a/src/nvmath/Fitting.h b/src/nvmath/Fitting.h
index 2421646..b2eeb15 100644
--- a/src/nvmath/Fitting.h
+++ b/src/nvmath/Fitting.h
@@ -13,7 +13,8 @@ namespace nv
 	void ComputeCovariance(int n, const Vector3 * points, const float * weights, Vector3::Arg metric, float * covariance);
 	Vector3 ComputePrincipalComponent(int n, const Vector3 * points, const float * weights, Vector3::Arg metric);
 
-	void Compute4Means(int n, const Vector3 * points, const float * weights, Vector3::Arg metric, Vector3 * cluster);
+	// Returns number of clusters [1-4].
+	int Compute4Means(int n, const Vector3 * points, const float * weights, Vector3::Arg metric, Vector3 * cluster);
 
 } // nv namespace