LeetCode / Remove Linked List Elements

https://leetcode.com/problems/remove-linked-list-elements/

Remove all elements from a linked list of integers that have value val.

Example:

Input: 1->2->6->3->4->5->6, val = 6
Output: 1->2->3->4->5

これがlistだったりsetだったら問題にもならないのですが、linked listならではのコードを書く必要があります。

解答・解説

誤答

各ノードを走査するための変数current_nodeを定義し、current_nodeを動かしながらvalを取り除いていくと考えます。

と、ここまでは良いのですが以下のようにやってしまうと誤答です。

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    def removeElements(self, head: ListNode, val: int) -> ListNode:
        
        current_node = head
        while current_node.next != None:
            if current_node.next.val == val:
                current_node.next = current_node.next.next
            else:
                current_node = current_node.next
                
        return head

これだと、1つ目の要素をチェックできないので
Inputが[1,2,6,3,4,5,6], val=6のときは良いですが
Inputが[1,2,6,3,4,5,6], val=1のときは正しい戻り値を返せません。

解法1

シンプルに1つ目の要素をチェックする処理を追加する、のが案1です。

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    def removeElements(self, head: ListNode, val: int) -> ListNode:
        
        while head is not None and head.val == val:
            head = head.next
        
        current = head
        while current is not None:
            if current.next is not None and current.next.val == val:
                current.next = current.next.next
            else:
                current = current.next
        
        return head

解法2

誤答のコードは1つ目の要素をチェックできないことが問題だったので、1つ目の要素の前にダミーのnodeを追加したlinked listを作ってやれば問題なくなる、と考えます。

そこで、dummy_headという変数を、ListNode(-1)とheadを結合させる形で作り、dummy_headの各nodeを走査します。

最後に返す値はdummy_head.nextである点に注意します（でないと、ダミーで追加したnodeも含めて返してしまう）。

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    def removeElements(self, head: ListNode, val: int) -> ListNode:
        
        dummy_head = ListNode(-1)
        dummy_head.next = head
        
        current_node = dummy_head
        while current_node.next != None:
            if current_node.next.val == val:
                current_node.next = current_node.next.next
            else:
                current_node = current_node.next
                
        return dummy_head.next