[Go/LeetCode] 21. Merge Two Sorted Lists

概要

↓の解法メモです。
https://leetcode.com/problems/merge-two-sorted-lists/description/

問題概要

二つのソートされたリストが渡されるので、これをマージします。
例えば [1,2,4] と [1,3,4] が渡されたとき、[1,1,2,3,4,4] を返します。

注意点として、ただのスライスが渡されるわけではありません。次の要素への参照を持った、リスト先頭のノード が渡されます。

 type ListNode struct {
     Val int
     Next *ListNode
 }

解法

初見

やることはシンプルで、リストの頭から小さいほうを取り出して新しいリストに入れるだけです。
そのままでは扱いづらいので、スライスよろしくインデックスのように参照できるよう byIndex 関数を実装しました。また、「配列の頭の要素を取り出す」「変数の参照先を次の要素に移す」操作をまとめる pop も実装しました。

func byIndex(list *ListNode, i int) *ListNode {
    if i == 0 {
        return list
    }
    if list.Next == nil {
        return nil
    }
    return byIndex(list.Next, i-1)
}
func pop(list *ListNode) (*ListNode, *ListNode) {
    p := list
    list = list.Next
    return p, list
}

/**
 * Definition for singly-linked list.
 * type ListNode struct {
 *     Val int
 *     Next *ListNode
 * }
 */
func mergeTwoLists(list1 *ListNode, list2 *ListNode) *ListNode {
    var mergedListNode *ListNode
    if list1 == nil && list2 == nil {
        return nil
    } else if  list1 != nil && list2 == nil {
        mergedListNode, list1 = pop(list1)
    } else if  list1 == nil && list2 != nil {
        mergedListNode, list2 = pop(list2)
    } else {
        if list1.Val < list2.Val {
            mergedListNode, list1 = pop(list1)
        } else {
            mergedListNode, list2 = pop(list2)
        }
    }

    for i:=0; list1 != nil || list2 != nil; i++ {
        if list1 != nil && list2 == nil {
            byIndex(mergedListNode, i).Next, list1 = pop(list1)
            continue
        } 
        if list1 == nil && list2 != nil {
            byIndex(mergedListNode, i).Next, list2 = pop(list2)
            continue
        }

        if list1.Val < list2.Val {
            byIndex(mergedListNode, i).Next, list1 = pop(list1)
            continue
        } else {
            byIndex(mergedListNode, i).Next, list2 = pop(list2)
            continue
        }
    }

    return mergedListNode
}

反省としては全く同じ条件式が2重に書かれている部分でしょうか。どうにかまとめたかったですが、条件式内部でやっている処理が別物なのでやりづらいな～と思いそのままにしています。
byIndex で毎回頭から要素を探索してるのもかなり無駄そう。最終的に返したいのが配列の先頭ノードなので mergedListNode が頭の参照をキープするようにしてましたが、もっといい書き方がありました。

最適化

Solutions から引用させてもらいます。
https://leetcode.com/problems/merge-two-sorted-lists/solutions/6251904/best-solution-ever-python-java-c-c-javascript-go-c-kotlin-typescript-swift

type ListNode struct {
    Val  int
    Next *ListNode
}

func mergeTwoLists(list1 *ListNode, list2 *ListNode) *ListNode {
    dummy := &ListNode{}
    op := dummy

    for list1 != nil && list2 != nil {
        if list1.Val <= list2.Val {
            op.Next = list1
            list1 = list1.Next
        } else {
            op.Next = list2
            list2 = list2.Next
        }
        op = op.Next
    }
    if list1 != nil {
        op.Next = list1
    }
    if list2 != nil {
        op.Next = list2
    }
    return dummy.Next
}

返す要素の頭への参照は dummy に持ちつつ、作業用のポインタを別に作って作業してます。全然これでよかった。
Go で構造体をリテラルから初期化するとゼロ値で初期化されてしまうので初期化処理を別に書いていましたが、別にポインタを持てばこれも解決しました。スッキリ。