2 files changed, 68 insertions, 27 deletions

diff --git a/src/line_accumulator.cc b/src/line_accumulator.cc
index bc85e6c..8865851 100644
--- a/src/line_accumulator.cc
+++ b/src/line_accumulator.cc
@@ -5,15 +5,23 @@
 
 namespace {
 
+std::uint8_t getRandomIndex(
+	std::mt19937&      random,
+	const std::uint8_t n
+) {
+	return std::uniform_int_distribution<std::uint8_t>{0, n}(random);
+}
+
 std::vector<std::uint8_t> getRandomIndizes(
-	std::random_device& device,
-	const std::uint8_t n) {
+	std::mt19937&      random,
+	const std::uint8_t n
+) {
 	std::vector<std::uint8_t> indizes(n);
 	std::iota(indizes.begin(), indizes.end(), 0);
 	std::shuffle(
 		indizes.begin(),
 		indizes.end(),
-		std::mt19937{device()}
+		random
 	);
 
 	return indizes;
@@ -24,6 +32,7 @@ std::vector<std::uint8_t> getRandomIndizes(
 LineAccumulator::LineAccumulator(const std::uint8_t max_length):
 	max_length_{max_length},
 	device_{},
+	random_{device_()},
 	length_{0},
 	tokens_{} { }
 
@@ -31,6 +40,10 @@ LineAccumulator::~LineAccumulator() {
 	this->discharge(false);
 }
 
+std::uint8_t LineAccumulator::getMissing() const {
+	return this->max_length_ - this->length_;
+}
+
 void LineAccumulator::operator()(const std::string& word) {
 	if ( ( this->length_ + word.length() ) > this->max_length_ ) {
 		this->discharge(true);
@@ -45,32 +58,56 @@ void LineAccumulator::operator()(const std::string& word) {
 	}
 }
 
+void LineAccumulator::justify() {
+	const std::uint8_t base = this->getMissing()
+	                        / (this->tokens_.size() - 1);
+
+	std::for_each(
+		this->tokens_.begin(),
+		this->tokens_.end()-2,
+		[&, base](auto& token) {
+			token.second  += base;
+			this->length_ += base;
+		}
+	);
+
+	switch ( this->getMissing() ) {
+		case 0: {
+			break;
+		}
+		case 1: {
+			this->tokens_[
+				getRandomIndex(this->random_, this->tokens_.size() - 1)
+			].second += 1;
+
+			break;
+		}
+		default: {
+			const auto indizes = getRandomIndizes(
+				this->random_,
+				this->tokens_.size() - 1
+			);
+
+			std::for_each(
+				indizes.begin(),
+				indizes.begin() + (this->max_length_ - this->length_),
+				[&](std::uint8_t x) {
+					this->tokens_[x].second += 1;
+					this->length_           += 1;
+				}
+			);
+
+			break;
+		}
+	}
+}
+
 void LineAccumulator::discharge(const bool full) {
+	this->length_              -= this->tokens_.back().second;
+	this->tokens_.back().second = 0;
+
 	if ( full ) {
-		this->length_              -= this->tokens_.back().second;
-		this->tokens_.back().second = 0;
-
-		const std::uint8_t missing = this->max_length_ - this->length_;
-		const std::uint8_t base    = missing / (this->tokens_.size()-1);
-
-		std::for_each(
-			this->tokens_.begin(),
-			this->tokens_.end()-2,
-			[base, this](auto& token) {
-				token.second  += base;
-				this->length_ += base;
-			}
-		);
-
-		auto indizes = getRandomIndizes(this->device_, this->tokens_.size()-1);
-
-		std::for_each(
-			indizes.begin(),
-			indizes.begin()+(this->max_length_ - this->length_),
-			[this](std::uint8_t x) {
-				this->tokens_[x].second += 1;
-			}
-		);
+		this->justify();
 	}
 
 	for ( const auto& token : this->tokens_ ) {
diff --git a/src/line_accumulator.h b/src/line_accumulator.h
index 93ebfc0..373cb1b 100644
--- a/src/line_accumulator.h
+++ b/src/line_accumulator.h
@@ -10,18 +10,22 @@ class LineAccumulator {
 		LineAccumulator(const std::uint8_t max_length);
 		~LineAccumulator();
 
+		std::uint8_t getMissing() const;
+
 		void operator()(const std::string& word);
 
 	private:
 		const std::uint8_t max_length_;
 
 		std::random_device device_;
+		std::mt19937       random_;
 		std::uint8_t       length_;
 
 		std::vector<
 			std::pair<std::string, std::uint8_t>
 		> tokens_;
 
+		void justify();
 		void discharge(const bool full);
 
 };