1 files changed, 23 insertions, 19 deletions

diff --git a/src/mass-graduation.js b/src/mass-graduation.js
index 46c4efe..69bbf07 100644
--- a/src/mass-graduation.js
+++ b/src/mass-graduation.js
@@ -33,27 +33,31 @@ function shuffleArray(array) {
 function chooseMandatory(graduations) {
 	graduations.sort((g1, g2) => g1.mandatory - g2.mandatory);
 
-	// Get all graduations with smaller mandatory amount than the one with the highest amount
-	let mandatoryCandidates = graduations.filter(
-		(g) => g.mandatory < graduations[graduations.length - 1].mandatory,
-	);
-	let rest;
-
-	// All graduations have the same mandatory amount
-	if (mandatoryCandidates.length === 0) {
-		mandatoryCandidates = graduations;
-		rest = [];
-	} else {
-		rest = graduations.filter((g) => !mandatoryCandidates.includes(g));
+	// Group the graduations according to the amount of mandatory changes
+	let mandatoryCandidateGroups = [];
+	let curGroup = [];
+	let mandatory = graduations[0].mandatory;
+	for (const g of graduations) {
+		if (g.mandatory === mandatory) {
+			curGroup.push(g);
+		} else {
+			// Start a new group
+			mandatoryCandidateGroups.push(curGroup);
+			curGroup = [g];
+			mandatory = g.mandatory;
+		}
 	}
+	// push the last group
+	mandatoryCandidateGroups.push(curGroup);
+
+	// Choose from the group with the smallest amount of mandatory changes
+	const chosen = [];
+	for (const candidateGroup of mandatoryCandidateGroups) {
+		if (chosen.length === 5) break;
 
-	// Get a random list of possible indexes
-	const candidatesIdxs = shuffleArray([...Array(mandatoryCandidates.length).keys()]);
-	let chosen = candidatesIdxs.slice(0, 5).map((i) => mandatoryCandidates[i]);
-	if (chosen.length < 5) {
-		const missing = 5 - chosen.length;
-		const restIdxs = shuffleArray([...Array(rest.length).keys()]);
-		chosen = chosen.concat(restIdxs.slice(0, missing).map((i) => rest[i]));
+		// Get a random list of possible indexes
+		const candidatesIdxs = shuffleArray([...Array(candidateGroup.length).keys()]);
+		chosen.push(...candidatesIdxs.slice(0, 5 - chosen.length).map((i) => candidateGroup[i]));
 	}
 
 	return chosen;