🐢

LeetCode - Squares of a Sorted Array をSwiftで解説する

5 min read
  • 元の問題は英語なので、代わりに翻訳したものを記載しています。
  • 回答はあくまでも例であり、他の方法も存在します。

前置き

Swiftで Sort アルゴリズム を書いたことがない人にとっては、勉強するのに良い問題でしょう。

問題

問題 難易度
Squares of a Sorted Array easy
整数の配列 nums が小さい順にソートされています。
各数値を2乗した配列を小さい順にソートして返します。

補足として、以下のような true を返すテストケースが与えられます。

Input: nums = [-4,-1,0,3,10]
Output: [0,1,9,16,100]

テストケースの計算過程はこうなります。

  1. 元の値   → [-4, -1, 0, 3, 10]
  2. 2乗する  → [16, 1, 0, 9, 100]
  3. ソートする → [0, 1, 9, 16, 100]

ポイントとしては、配列にマイナスの値がある場合、2乗した際に値がプラスに反転します。その際、最初は小さい順にソートされていた配列の順番が変わってしまうため、再度ソートする必要があります。

答え

  • Swift のAPIを使用する方法
  • Sort アルゴリズムを用いる方法

の2つを紹介します。

1. Swift のAPIを使用する方法

func sortedSquares(_ A: [Int]) -> [Int] {
    return A.map { $0 * $0 }.sorted()
}

Swiftでは sorted メソッドが用意されています。

sorted メソッドは、内部的には複数の Sort アルゴリズムを組み合わせています。そのため自前で Sort のアルゴリズムを実装するよりも、こちらを使用する方が高速でしょう。
(参考: Swiftではどのようにソートを行っているか

しかし、今回はあえてアルゴリズムを使った方法を見ていきます。

2. Sort アルゴリズムを用いる方法

ソートのアルゴリズムはいくつか存在しますが、今回は2種類を紹介します。

① Bubble Sort(バブルソート)

隣り合う要素の値を比較し、条件に応じて交換するアルゴリズムです。

bubble sort
(引用:バブルソートでシンプルなベンチマーク)

では、実際のコードを見ていきましょう。

func sortedSquares(_ A: [Int]) -> [Int] {
    // ① 2乗した配列を作る
    var result = [Int]()
    for n in A {
        result.append(n*n)
    }
    
    // ② ソートしていく
    var isSwapped = true
    while isSwapped {
        // 一度でも値の交換をした場合は、再度比較をし直すためのフラグ
        isSwapped = false
        
        // 最初から数字を比較していく
        // 中の計算で'i-1'をしても問題ないように'0'ではなく'1'から始まっている
        for i in 1..<result.count {
            // 隣り合う数字の比較
            guard result[i] < result[i-1] else { continue }
            // 該当する場合は'swapAt'で値を交換する
            result.swapAt(i, i-1)
            isSwapped = true
        }
    }
    return result
}

再起関数使う書き方もありますが、今回はwhileフラグを使って書いています。

今回こちらを紹介しましたが、Bubble Sort は総当たりで比較していくため、計算量が非常に多く時間がかかります。この問題で使用すると、残念ながらタイムアウトになってしまいます。

では、こちらよりも高速な Sort アルゴリズムを見てみましょう。

② Merge Sort(マージソート)

ソートされていないリストを2つのリストに分割し、それぞれをソートさせた後、それらを1つにする(マージする)アルゴリズムです。

bubble sort
(引用:ソースコード探検隊 » アルゴリズムとデータ構造 » アルゴリズム » マージソート)

では、実際のコードを見ていきましょう。

func sortedSquares(_ A: [Int]) -> [Int] {
    // ① 2乗した配列を作る
    var squaredNums = [Int]()
    for n in A {
        squaredNums.append(n*n)
    }
    
    // ② ソートしていく & ③ 合成する 
    return mergeSort(squaredNums)
}

func mergeSort(_ lists: [Int]) -> [Int] {
    guard 1 < lists.count else { return lists }
    
    // ②-1 配列を2つに分割
    let count = lists.count
    let left = Array(lists[..<Int(count/2)])
    
    // ②-2 分割した個々の配列をソートし続ける
    let right = Array(lists[count/2...count-1])
    return merge(left: mergeSort(left), right: mergeSort(right))
}

func merge(left: [Int], right: [Int]) -> [Int] {
    var left = left, right = right, result = [Int]()
    
    // ③-1 2つの配列から最初の値を見て、小さい値から順にリストに追加していくことで、2つの配列を合成していく
    while 0 < left.count && 0 < right.count {
        // `removeFirst`で追加した方の先頭の値を除外していく
        if left[0] <= right[0] {
            result.append(left.removeFirst())
        } else {
            result.append(right.removeFirst())
        }
    }
    
    // ③-2 `while`ではいずれかが'0'になった時に処理が終わってしまうので、最後の値を配列に追加する
    result.append(contentsOf: left)
    result.append(contentsOf: right)
    
    return result
}

Merge Sort では「 sort する関数」と「 merge する関数」をそれぞれ用意します。一見するとコード量は多いですが、Bubble Sort とは違い、個々の配列を見ることで比較する値の計算量が減少されます。

2つの Sort を紹介しましたが、他にもアルゴリズムはいくつかあるので、興味がある方は調べてみてください。

実行結果

Kind Runtime Memory
sorted() 380ms~ 16MB前後
Bubble Sort Time Limit Exceeded -
Merge Sort 740ms~ 16MB前後

引用

Discussion

ログインするとコメントできます