10 lines, 24 ms, iterative, one-pass c++ code with explanation


  • 0
    P
    class Solution {
    public:
    
        template<typename T> void shift(T& a, T& b, T& c) {
            T t = a; a = b; b = c; c = t;
        }
        
        ListNode** reverse(ListNode** head, int k) {
            ListNode *p = *head, *next = nullptr;
            
            while (p && k--) shift(p->next, next, p);
    
            ListNode** res = &(*head)->next;
            *res = p;
            *head = next;
            return k > 0 ? reverse(head, -1) : res;
        }
        
        ListNode* reverseKGroup(ListNode* head, int k) {
            for (ListNode** pp = &head; *pp; pp = reverse(pp, k));
            return head;
        }
    };
    

    The code looks recursive, but it is not. The basic idea is to count the number of nodes in each group. If the number is less than k then we should reverse it again. This ensures that every k-group is reversed only once while the last group of remaining nodes is reversed twice. Since we are expecting multiple k-groups and only one group of remaining nodes, the code is almost one-pass.

    The core function ListNode** reverse(ListNode** head, int k) accepts a ListNode** that corresponds to the head of the current group, and returns ListNode** corresponding to the head of the next group. A helper ListNode* next is used to store the new next field of each reversed node, and a ListNode* p points to the node whose next field is being reversed. Each reverse is then a shift of p->next, next and p.

    For example, let k = 3:

    ---> 1 ---> 2 ---> 3 ---> ?
     ^
     |
    head
    

    After a round of shift:

        nil
         ^
         |
    ---> 1      2 ---> 3 ---> ?
     ^   ^      ^
     |   |      |
     |  next    p
    head
    

    After k rounds of shift:

        nil
         ^
         |
    ---> 1 <--- 2 <--- 3      ?
     ^                 ^      ^
     |                 |      |
     |                next    p
    head
    

    The rest of the work is simply setting up *head, the next field of the original group head and the return value.

    For the last group with less than k nodes, we invoke reverse() again, but pass -1 as k, which leads to another reversion till the end. Again, calling reverse() inside reverse() will happen ONLY ONCE.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.