Partition List Editorial

Practice Problem Link: Partition List

Please make sure to try solving the problem yourself before looking at the editorial.

Problem Statement

Given a linked list and a key, partition it such that all nodes with values less than key are before it and all nodes with values greater than key are after it.

The relative order of the numbers in each of the partition should be maintained.

Approach

• Create three lists to store elements that are less than, equal to and greater than the key.
• Now traverse the given array and insert each node into its respective list in the same order.
• Merge the three linked lists and return its head.

Analysis

• Time Complexity: O(n)
• Auxiliary Space Complexity: O(1)

Implementation

C++

/* This is the ListNode class definition

class ListNode {
public:
	int data;
	ListNode* next;

	ListNode(int data) {
		this->data = data;
		this->next = NULL;
	}
};
*/
ListNode* partitionList(ListNode* head, int key) {
	ListNode* smallerHead = NULL, *equalHead = NULL, *greaterHead = NULL;
	ListNode* smallerLast = NULL, *equalLast = NULL, *greaterLast = NULL;
	while (head != NULL) {
		if (head->data < key) {
			if(smallerHead == NULL) {
				smallerHead = smallerLast = head;
			} else {
				smallerLast->next = head;
				smallerLast = smallerLast->next;
			} 
		} else if (head->data == key) {
			if (equalHead == NULL) {
				equalHead = equalLast = head;
			} else {
				equalLast->next = head;
				equalLast = equalLast->next;
			}
		} else {
			if (greaterHead == NULL) {
				greaterHead = greaterLast = head;
			} else {
				greaterLast->next = head;
				greaterLast = greaterLast->next;
			}
		}
		head = head->next;
	}
	if (greaterHead != NULL) {
		greaterLast->next = NULL;
	}
	if (equalHead == NULL) {
		if (smallerHead == NULL) {
			smallerHead = greaterHead;
		} else {
			smallerLast->next = greaterHead;
		}
	} else {
		equalLast->next = greaterHead;
		if (smallerHead == NULL) {
			smallerHead = equalHead;
		} else {
			smallerLast->next = equalHead;
		}
	}
	return smallerHead;
}

Java

/** This is the ListNode class definition

class ListNode {
	int data;
	ListNode next;

	ListNode(int data) {
		this.data = data;
		this.next = null;
	}
}
**/
class Solution {
	ListNode partitionList(ListNode head, int key) {
		ListNode smallerHead = null, equalHead = null, greaterHead = null;
		ListNode smallerLast = null, equalLast = null, greaterLast = null;
		while (head != null) {
			if (head.data < key) {
				if(smallerHead == null) {
					smallerHead = smallerLast = head;
				} else {
					smallerLast.next = head;
					smallerLast = smallerLast.next;
				} 
			} else if (head.data == key) {
				if (equalHead == null) {
					equalHead = equalLast = head;
				} else {
					equalLast.next = head;
					equalLast = equalLast.next;
				}
			} else {
				if (greaterHead == null) {
					greaterHead = greaterLast = head;
				} else {
					greaterLast.next = head;
					greaterLast = greaterLast.next;
				}
			}
			head = head.next;
		}
		if (greaterHead != null) {
			greaterLast.next = null;
		}
		if (equalHead == null) {
			if (smallerHead == null) {
				smallerHead = greaterHead;
			} else {
				smallerLast.next = greaterHead;
			}
		} else {
			equalLast.next = greaterHead;
			if (smallerHead == null) {
				smallerHead = equalHead;
			} else {
				smallerLast.next = equalHead;
			}
		}
		return smallerHead;
	}
}