How to partition a linked list

Question 2.4 of Cracking the Coding Interview:

Write code to partition a linked list around a value x, such that all nodes less than x come before all nodes greater than or equal to x.

My first approach built two new lists: one for values less than x, and another for values greater than or equal to x. When we’re done, I returned the concatenation of these lists. To concatenate the two lists efficiently, I remembered the tail of the first list.

Then I read the answer, and found a simpler implementation: you only have to build one new list! When you find a value less than x, you put the node at the head of the list; otherwise, you append the node at the tail. To make append constant-time, you keep a pointer to the tail. This runs in O(length of input), which is optimal.

Here’s an implementation in C:

#include <stdlib.h>
#include <assert.h>
#include <stdarg.h>

typedef struct Node {
  int val;
  struct Node * next;
} Node;

Node* partition(int x, Node* list) {
  if (list == NULL) return NULL;
  Node* head = list;
  Node* tail = list;
  Node* n = list->next;
  list->next = NULL;
  while (n != NULL) {
    Node* next = n->next;
    if (n->val < x) {
      n->next = head;
      head = n;
    } else {
      n->next = NULL;
      tail->next = n;
      tail = n;
    }
    n = next;
  }

  return head;
}

// ----------------------------------------------------
// --------------------- TESTS ------------------------

Node* mk_node(int val, Node* next) {
  Node* node = malloc(sizeof(Node));
  node->val = val;
  node->next = next;
  return node;
}

void assert_lists_eq(Node* actual, Node* expected) {
  while (actual != NULL) {
    assert(expected != NULL);
    assert(actual->val == expected->val);
    actual = actual->next;
    expected = expected->next;
  }
  assert(expected == NULL);
}

Node* mk_list(int len, ...) {
  Node* list_start = NULL;
  Node* list_end = NULL;
  va_list argp;
  va_start(argp, len);
  for (int i = 0; i < len; i++) {
    Node* node = mk_node(va_arg(argp, int), NULL);
    if (list_start == NULL) {
      list_start = node;
      list_end = node;
    } else {
      list_end->next = node;
      list_end = node;
    }
  }
  va_end(argp);
  return list_start;
}

int main(int argc, char** argv) {
  assert_lists_eq(
    partition(5,  mk_list(7,    2,6,3,7,5,3,6)), 
                  mk_list(7,    3,3,2,6,7,5,6));

  assert_lists_eq(
    partition(1,  mk_list(7,    2,6,3,7,5,3,6)), 
                  mk_list(7,    2,6,3,7,5,3,6));

  assert_lists_eq(
    partition(42, mk_list(7,    2,6,3,7,5,3,6)), 
                  mk_list(7,    6,3,5,7,3,6,2));

  assert_lists_eq(
    partition(5,  mk_list(0)), 
                  mk_list(0));

  return 0;
}
Tagged #ctci, #programming, #c.

Similar posts

More by Jim

Want to build a fantastic product using LLMs? I work at Granola where we're building the future IDE for knowledge work. Come and work with us! Read more or get in touch!

This page copyright James Fisher 2020. Content is not associated with my employer. Found an error? Edit this page.