rustgym/leetcode/
_25_reverse_nodes_in_k_group.rs

1struct Solution;
2use rustgym_util::*;
3use std::collections::VecDeque;
4
5impl Solution {
6    fn reverse_k_group(head: ListLink, k: i32) -> ListLink {
7        let mut p = head;
8        let mut count = 0;
9        let mut queue: VecDeque<ListLink> = VecDeque::new();
10        let k = k as usize;
11        while let Some(mut node) = p {
12            p = node.next.take();
13            queue.push_back(Some(node));
14            count += 1;
15            if count == k {
16                break;
17            }
18        }
19        if queue.len() == k {
20            let mut prev: ListLink = Self::reverse_k_group(p, k as i32);
21            while let Some(link) = queue.pop_front() {
22                if let Some(mut node) = link {
23                    node.next = prev;
24                    prev = Some(node);
25                }
26            }
27            prev
28        } else {
29            let mut prev: ListLink = None;
30            while let Some(link) = queue.pop_back() {
31                if let Some(mut node) = link {
32                    node.next = prev;
33                    prev = Some(node);
34                }
35            }
36            prev
37        }
38    }
39}
40
41#[test]
42fn test() {
43    let head = list!(1, 2, 3, 4, 5);
44    let k = 2;
45    let res = list!(2, 1, 4, 3, 5);
46    assert_eq!(Solution::reverse_k_group(head, k), res);
47    let head = list!(1, 2, 3, 4, 5);
48    let k = 3;
49    let res = list!(3, 2, 1, 4, 5);
50    assert_eq!(Solution::reverse_k_group(head, k), res);
51    let head = list!(1, 2);
52    let k = 2;
53    let res = list!(2, 1);
54    assert_eq!(Solution::reverse_k_group(head, k), res);
55}