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}