Search code examples
c++recursiontreesingly-linked-list

flatten a singly linked list

I have this tree structure, where each node has a next and child pointer. I am trying to flatten such a tree to a linked list where each node only uses the next pointer.

Example input:

1 --> 2 --> 3 --> 4 --> 5 --> 6
            |
            V 
            7 --> 8 -- > 9 --> 10
                  |
                  V
                  11 --> 12

Expected output:

  1 -> 2 -> 3 -> 4 -> 5 -> 6 -> 7 -> 8 -> 9 -> 10 -> 11 -> 12

This is my code:

struct ListNode {
    int val;
    struct ListNode* next;
    struct ListNode* child;

    ListNode(int x) {
        val = x;
        next = NULL;
        child = NULL;
    }
};

ListNode* flatten(ListNode* head) {
    ListNode *temp = head, *r = NULL;
    while (temp) {
        if (temp->child) {
            r = temp->next;
            temp->next = flatten(temp->child);
        }
        if (!temp->next && r) {
            temp->next = r;
            r = NULL;
        }
        temp = temp->next;
    }
    return head;
}

Problem

When I provide it the example input, my code gets into infinite recursion.

Where is my mistake?

Solution

  • Your code does not get into infinite recursion, but into an infinite loop.

    There are these issues:

    • The child pointers are never cleared to nullptr, and so you get nodes that have both their child and next pointing to the same node. As you also link back tail nodes to nodes in upper "levels", this eventually leads to revisits of the same child nodes over and over again. Things would not get into such cycles if you would clear child pointers with temp->child = nullptr; right after the recursive call comes back.

    • Although r saves a pointer to be reused later, there is no provision for when multiple nodes in the same "horizontal" list have a non-null child pointer: your loop can only remember one of them in r, and so you will lose information.

    • The logic gives precedence to linking child lists, coming back to the "calling" list after child lists have been wired into the current list. Yet your desired output gives precedence to the next lists, wanting to put child nodes after those. I'll assume this part of the code is wrong, and your description of the expected output is right.

    You can do this without recursion by just keeping track of the current tail of the part of the list that is being flattened, while the temporary pointer walks behind it to move and clean up the child pointers it encounters on its way.

    ListNode* getTail(ListNode* head) { // Assumes non-null argument
    while (head->next) {
        head = head->next;
    }
    return head;
}

ListNode* flatten(ListNode* head) {
    ListNode *tail = getTail(head);
    for (ListNode *node = head; node; node = node->next) {
        tail->next = node->child;
        tail = getTail(tail);
        node->child = nullptr;
    }
    return head;
}

    Unrelated to your question, but in C++:

    So:

    struct ListNode {
    int val;
    ListNode* next = nullptr;
    ListNode* child = nullptr;

    ListNode(int x): val(x) {}
}

    Finally, flatten will always return the argument given to it, so you could consider making it a void function.