Browse Source

feat: partition_list

Zhang Li 1 year ago
parent
commit
fec7908e25
2 changed files with 43 additions and 0 deletions
  1. 1 0
      src/solutions/list/mod.rs
  2. 42 0
      src/solutions/list/partition_list.rs

+ 1 - 0
src/solutions/list/mod.rs

@@ -3,3 +3,4 @@ mod rotate_list;
 mod remove_duplicate_from_sorted_list;
 mod merge_k_sorted_lists;
 mod merge_two_sorted_lists;
+mod partition_list;

+ 42 - 0
src/solutions/list/partition_list.rs

@@ -0,0 +1,42 @@
+struct Solution;
+use crate::solutions::ListNode;
+
+impl Solution {
+    pub fn partition(head: Option<Box<ListNode>>, x: i32) -> Option<Box<ListNode>> {
+        let mut head = head;
+        let mut before: Option<Box<ListNode>> = None;
+        let mut after: Option<Box<ListNode>> = None;
+        let mut before_tail = &mut before;
+        let mut after_tail = &mut after;
+        while let Some(mut node) = head {
+            head = node.next.take();
+            if node.val < x {
+                *before_tail = Some(node);
+                before_tail = &mut before_tail.as_mut().unwrap().next
+            } else {
+                *after_tail = Some(node);
+                after_tail = &mut after_tail.as_mut().unwrap().next
+            }
+
+        }
+        *before_tail = after;
+        before
+
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::solutions::{to_list, to_vec};
+
+    #[test]
+    fn test_partition() {
+        let head = vec![1,4,3,2,5,2];
+
+        let head = to_list(head);
+        let x = 3;
+        let res = Solution::partition(head, x);
+        assert_eq!(to_vec(res), vec![1,2,2,4,3,5])
+    }
+}