Merge pull request #43 from Juice-jl/issue#42

fix bug for k-means in em
Tractables · Dec 3, 2020 · 6340103 · 6340103
2 parents ec771f5 + 109f6e3
commit 6340103
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/mixtures/em.jl b/src/mixtures/em.jl
@@ -52,13 +52,13 @@ function clustering(data, mix_num::Int64; maxiter=200)::Vector
         return [data]
     end
     data = Matrix(data)
-    R = kmeans(data, mix_num; maxiter=maxiter)
+    R = kmeans(data', mix_num; maxiter=maxiter)
     @assert nclusters(R) == mix_num
     a = assignments(R)
 
     clustered_data = Vector()
     for k in 1 : mix_num
-        push!(clustered_data, DataFrame(data[:, findall(x -> x == k, a)]'))
+        push!(clustered_data, DataFrame(data[findall(x -> x == k, a), :]))
     end
 
     return clustered_data