分隔链表（LeetCode）

题目

解题

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


def partition(head: ListNode, x: int) -> ListNode:
    # 创建两个虚拟头节点
    smaller_head = ListNode(0)
    greater_head = ListNode(0)

    # 用两个指针来操作这两个链表
    smaller = smaller_head
    greater = greater_head

    # 遍历原链表
    while head:
        next_node = head.next  # 先保存下一个节点
        if head.val < x:
            # 如果节点值小于 x，将其加入到小链表
            smaller.next = head
            smaller = smaller.next
        else:
            # 否则加入到大链表
            greater.next = head
            greater = greater.next

        head.next = None  # 断开当前节点的 next 指针
        head = next_node  # 移动到下一个节点

    # 将大链表的末尾指向 None
    greater.next = None
    # 将小链表的末尾连接到大链表的头部
    smaller.next = greater_head.next

    return smaller_head.next


def create_linked_list(values):
    # 通过值列表创建链表
    if not values:
        return None
    head = ListNode(values[0])
    current = head
    for value in values[1:]:
        current.next = ListNode(value)
        current = current.next
    return head


def print_linked_list(head: ListNode):
    # 打印链表的值
    values = []
    while head:
        values.append(head.val)
        head = head.next
    print(" -> ".join(map(str, values)))


if __name__ == '__main__':
    test_values = [-1, 1, 2, 2, 4, 3, 5]
    test_x = 3
    # 创建链表
    test_head = create_linked_list(test_values)

    # 分隔链表
    new_head = partition(test_head, test_x)

    # 打印结果
    print("分隔后的链表:")
    print_linked_list(new_head)