diff --git a/parcels/collection/collectionaos.py b/parcels/collection/collectionaos.py
index adddb8fe7..7917906d0 100644
--- a/parcels/collection/collectionaos.py
+++ b/parcels/collection/collectionaos.py
@@ -3,6 +3,7 @@
 
 from ctypes import c_void_p
 
+import copy
 import numpy as np
 
 from parcels.collection.collections import ParticleCollection
@@ -103,7 +104,17 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, p
                         if mpi_rank == 0:
                             coords = np.vstack((lon, lat)).transpose()
                             kmeans = KMeans(n_clusters=mpi_size, random_state=0).fit(coords)
-                            self._pu_indicators = kmeans.labels_
+                            labels = np.linspace(0, mpi_size, lon.size, endpoint=False)
+                            labels = np.floor(labels)
+                            reorderedCoords = copy.copy(labels)
+                            for i in range(0, len(kmeans.cluster_centers_)):
+                                clusterCentre = kmeans.cluster_centers_[i]
+                                distances = map(lambda point: ((point[0]-clusterCentre[0])**2 + (point[1]-clusterCentre[1])**2), coords)
+                                sortedDistanceIdxs = np.argsort(list(distances))
+                                numberToChoose = sum(labels == i)
+                                reorderedCoords[sortedDistanceIdxs[0:numberToChoose]] = i
+                                coords[sortedDistanceIdxs[0:numberToChoose],:] = float('inf')
+                            self._pu_indicators = reorderedCoords
                         else:
                             self._pu_indicators = None
                         self._pu_indicators = mpi_comm.bcast(self._pu_indicators, root=0)
diff --git a/parcels/collection/collectionsoa.py b/parcels/collection/collectionsoa.py
index f739acad8..ba84cd08b 100644
--- a/parcels/collection/collectionsoa.py
+++ b/parcels/collection/collectionsoa.py
@@ -3,7 +3,7 @@
 from ctypes import Structure, POINTER
 from bisect import bisect_left
 from math import floor
-
+import copy
 import numpy as np
 
 from parcels.collection.collections import ParticleCollection
@@ -101,10 +101,20 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, p
             if mpi_size > 1:
                 if partitions is not False:
                     if (self._pu_indicators is None) or (len(self._pu_indicators) != len(lon)):
-                        if mpi_rank == 0:
+                        if mpi_rank == 0:                            
                             coords = np.vstack((lon, lat)).transpose()
                             kmeans = KMeans(n_clusters=mpi_size, random_state=0).fit(coords)
-                            self._pu_indicators = kmeans.labels_
+                            labels = np.linspace(0, mpi_size, lon.size, endpoint=False)
+                            labels = np.floor(labels)
+                            reorderedCoords = copy.copy(labels)
+                            for i in range(0, len(kmeans.cluster_centers_)):
+                                clusterCentre = kmeans.cluster_centers_[i]
+                                distances = map(lambda point: ((point[0]-clusterCentre[0])**2 + (point[1]-clusterCentre[1])**2), coords)
+                                sortedDistanceIdxs = np.argsort(list(distances))
+                                numberToChoose = sum(labels == i)
+                                reorderedCoords[sortedDistanceIdxs[0:numberToChoose]] = i
+                                coords[sortedDistanceIdxs[0:numberToChoose],:] = float('inf')
+                            self._pu_indicators = reorderedCoords
                         else:
                             self._pu_indicators = None
                         self._pu_indicators = mpi_comm.bcast(self._pu_indicators, root=0)