n0025_reverse_nodes_in_k_group.rs

/**
 * [25] Reverse Nodes in k-Group
 *
 * Given a linked list, reverse the nodes of a linked list k at a time and return its modified list.
 * 
 * k is a positive integer and is less than or equal to the length of the linked list. If the number of nodes is not a multiple of k then left-out nodes in the end should remain as it is.
 * 
 * 
 * 
 * 
 * Example:
 * 
 * Given this linked list: 1->2->3->4->5
 * 
 * For k = 2, you should return: 2->1->4->3->5
 * 
 * For k = 3, you should return: 3->2->1->4->5
 * 
 * Note:
 * 
 * 
 * 	Only constant extra memory is allowed.
 * 	You may not alter the values in the list's nodes, only nodes itself may be changed.
 * 
 * 
 */
pub struct Solution {}
use super::util::linked_list::{ListNode, to_list};

// submission codes start here

impl Solution {
    pub fn reverse_k_group(head: Option<Box<ListNode>>, k: i32) -> Option<Box<ListNode>> {
        let mut dummy_head = Some(Box::new(ListNode{val: 0, next: head}));
        let mut head = dummy_head.as_mut();
        'outer: loop {
            let mut start = head.as_mut().unwrap().next.take();
            if start.is_none() { break 'outer }
            let mut end = start.as_mut();
            for _ in 0..(k-1) {
                end = end.unwrap().next.as_mut();
                if end.is_none() {
                    head.as_mut().unwrap().next = start;
                    break 'outer
                }
            }
            let mut tail = end.as_mut().unwrap().next.take();
            // BEFORE: head -> start -> 123456... -> end   -> tail
            // AFTER:  head -> end   -> ...654321 -> start -> tail
            let end = Solution::reverse(start, tail);
            head.as_mut().unwrap().next = end;
            for _ in 0..k {
                head = head.unwrap().next.as_mut()
            }
        }
        dummy_head.unwrap().next
    }

    #[inline(always)]
    fn reverse(mut head: Option<Box<ListNode>>, tail: Option<Box<ListNode>>) -> Option<Box<ListNode>> {
        let mut prev = tail;
        let mut current = head;
        while let Some(mut current_node_inner) = current {
            let mut next = current_node_inner.next.take();
            current_node_inner.next = prev.take();
            prev = Some(current_node_inner);
            current = next;
        }
        prev
    }
}

// submission codes end

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_25() {
        assert_eq!(Solution::reverse_k_group(to_list(vec![1,2,3,4,5]), 2),
                   to_list(vec![2,1,4,3,5]));
        assert_eq!(Solution::reverse_k_group(to_list(vec![1,2,3,4,5]), 3),
                   to_list(vec![3,2,1,4,5]));
        assert_eq!(Solution::reverse_k_group(to_list(vec![1,2,3,4,5]), 5),
                   to_list(vec![5,4,3,2,1]));
        assert_eq!(Solution::reverse_k_group(to_list(vec![1]), 1),
                   to_list(vec![1]));
    }
}