Function partition

pub fn partition(a: impl AsRef<Array>, kth: i32) -> Result<Array>

Returns a partitioned copy of the flattened array such that the smaller kth elements are first. Returns an error if the arguments are invalid.

The ordering of the elements in partitions is undefined.

Params

• array: input array
• kth: Element at the kth index will be in its sorted position in the output. All elements before the kth index will be less or equal to the kth element and all elements after will be greater or equal to the kth element in the output.

Example

use mlx_rs::{Array, ops::*};

let a = Array::from_slice(&[3, 2, 1], &[3]);
let kth = 1;
let result = partition(&a, kth);