modify elbow.py for better n_clusters

Author: 东凡

elbow.py

@@ -5,6 +5,9 @@
 
 
 def get_distance(data):
+    """
+    把数据聚成多个类别，输出各个类别下的组内方差和，画出"肘型图"
+    """
     distance = []
     k = []
     for n_clusters in range(1, 10):
@@ -27,17 +30,24 @@ def get_distance(data):
 
 
 def get_elbow(distance):
+    """
+    输出最佳聚类个数
+    """
     a_1 = pd.Series(distance[1:])
     a_2 = pd.Series(distance[:-1])
 
     tmp = a_2 / a_1
-    tmp[tmp < 1] = np.inf
-    return tmp.values.argmax() + 2
+
+    tmp_1 = pd.Series(tmp[1:].values)
+    tmp_2 = pd.Series(tmp[:-1].values)
+
+    tmp2 = tmp_2 / tmp_1
+    return tmp2.values.argmax() + 2
 
 
 if __name__ == '__main__':
     a = pd.DataFrame(np.array([1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4]) * 10 + np.random.rand(15))
-    distance, k = get_distance(grey)
+    distance, k = get_distance(a)
     best_n_cluster = get_elbow(distance)
     best_n_cluster