From 04d28f2a74802ffa92a5ff9c8845d7b4db8ee3cc Mon Sep 17 00:00:00 2001
From: encyclopaedia <prashanth-hegde@users.noreply.github.com>
Date: Sat, 23 Sep 2023 07:52:54 -0500
Subject: [PATCH] arrays: add a partition function, that splits a given array,
 based on a criteria, passed as a callback fn (#19417)

---
 vlib/arrays/arrays.v      | 15 +++++++++++++++
 vlib/arrays/arrays_test.v | 17 +++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/vlib/arrays/arrays.v b/vlib/arrays/arrays.v
index 8e9302e8fe455f..d128640aa609e9 100644
--- a/vlib/arrays/arrays.v
+++ b/vlib/arrays/arrays.v
@@ -728,3 +728,18 @@ pub fn join_to_string[T](array []T, separator string, transform fn (elem T) stri
 	}
 	return sb.str()
 }
+
+// partition splits the original array into pair of lists,
+// where first list contains elements for which predicate yielded true,
+// while second list contains elements for which predicate yielded false
+pub fn partition[T](array []T, predicate fn (elem T) bool) ([]T, []T) {
+	mut matching, mut non_matching := []T{}, []T{}
+	for item in array {
+		if predicate(item) {
+			matching << item
+		} else {
+			non_matching << item
+		}
+	}
+	return matching, non_matching
+}
diff --git a/vlib/arrays/arrays_test.v b/vlib/arrays/arrays_test.v
index f2dfba55754fb5..55e7edd947823d 100644
--- a/vlib/arrays/arrays_test.v
+++ b/vlib/arrays/arrays_test.v
@@ -481,3 +481,20 @@ fn test_join_to_string() {
 		return '1'
 	}) == ''
 }
+
+fn test_partition() {
+	a := [1, 2, 3, 4, 5, 6, 7, 8]
+	lower, upper := partition(a, fn (it int) bool {
+		return it < 5
+	})
+	assert lower.len == 4
+	assert upper.len == 4
+	assert lower == [1, 2, 3, 4]
+	assert upper == [5, 6, 7, 8]
+
+	lower2, upper2 := partition(a, fn (it int) bool {
+		return it < 1
+	})
+	assert lower2.len == 0
+	assert upper2.len == 8
+}