HackerRankのSherlock and Arrayを解く

HackerRankの問題をRubyで解いた解答になります。
今回は、Sherlock and Arrayに挑戦しました。

問題内容

問題

WatsonはSherlockに整数の配列を与えます。彼の課題は、左側のすべての要素の合計が右側のすべての要素の合計と等しくなるような配列の要素を見つけることです。たとえば、arr = [5,6,8,11] が与えられた場合、8は合計が11になる2つのサブ配列の間にあります。もし与えらた配列が[1]の場合、左右それぞれの合計は0になるというルールがあります。

あなたは与えられた整数の配列が、ある要素の左右のそれぞれの合計が等しくなるという基準を満たす要素があるかどうかを確認する必要があります。

関数の説明

入力フォーマット

最初の行Tには、テストケースの数が含まれています。
次の2行Tはそれぞれテストケースを表しています。

• 最初の行nは、配列arrの要素数になります。
• 2行目のnには、0 ≦ i < nの範囲でスペースで区切られた整数arr[i]が含まれています。

制約

• 1 ≦ T ≦ 10
• 1 ≦ T ≦ 105
• 1 ≦ arr[i] ≦ 2 * 104
• 0 ≦ i < n

出力フォーマット

各テストケースで、配列に要素が存在する場合は、YESを出力します。つまり、左側の要素の合計が右側の要素の合計と等しくなります。それ以外の場合はNOを出力します。

サンプルインプット

2
3
1 2 3
4
1 2 3 3

サンプルアウトプット

NO
YES

解答

ここからは、僕の解答になります。
いきなり正解が書けたわけではなく、順を追って徐々に改良していき解答にこぎつけた形です。先に正解と判断されたコードだけのせておきます。

他にもたくさんの書き方があると思うので、質問やより良い書き方があれば、TwitterやGithubからつっこみを入れてください。

def balancedSums(arr)
    n = arr.size
    x = 0
    sum = arr.sum
    if n == 1
        return 'YES'
    else
        arr.each do |y|
            if 2 * x == sum - y
                return 'YES'
            end
            x += y
        end
        return 'NO'
    end
end

最初のコード

最初に僕が書いたコードはこんな感じでした。
これだと、サンプルは成功しますが、要素が1つの時満たせなかったり、要素の数が増えるとタイムアウトになってしまうなど問題がありました。

def balancedSums(arr)
    n = arr.size
    i = n - 2
    right_all = arr.slice(1..-1).sum
    if right_all == 0
        result = 'YES'
    else
        i.times do |j|
            if j == i
                left_all = arr.slice(1..-1).sum
                if left_all == 0
                    result = 'YES'
                    break
                else
                    result = 'NO'
                end
            else
                arr_left = arr.slice(0..j)
                arr_right =arr.slice(j+2..-1)
                if arr_left.sum == arr_right.sum
                    result = 'YES'
                    break
                else
                    result = 'NO'
                end
            end
        end
    end
    return result
end

ループの中身を変更してみる

まずは、ループ内をどうしたら改善できるか（スピードを上げられるか）を考え、改良を試みました。境となる要素が最初の時と最後の時を先に確認してしまって、それ以外を計算するなど色々試行錯誤してみましたが、特に結果は変わりませんでした。

def balancedSums(arr)
    n = arr.size
    i = n - 2
    right_all = arr.slice(1..-1).sum
    left_all = arr.slice(0..-2).sum
    #p right_all
    #p left_all
    if right_all == 0
        result = 'YES'
    elsif left_all == 0
        result = 'YES'
    else
        i.times do |j|
            if j == i
                left_all = arr.slice(1..-1).sum
                if left_all == 0
                    result = 'YES'
                    break
                else
                    result = 'NO'
                end
            arr_left = arr.slice(0..j).sum
            arr_right = arr.slice(j+2..-1).sum
            if arr_left == arr_right
                result = 'YES'
                break
            else
                arr_left = arr.slice(0..j)
                arr_right =arr.slice(j+2..-1)
                if arr_left.sum == arr_right.sum
                    result = 'YES'
                    break
                else
                    result = 'NO'
                end
                result = 'NO'
            end
        end
    end
    return result
end

Discussionを見てみる

配列同士の差集合のみの合計を比較できないか？など考えたりもしましたが、いい方法が思い浮かばず、違う考え方を取り入れるためにもタイムアウトに関してどんな議論が行われているのかを「Discussion」の部分で見てみることにしました。

同じようにタイムアウトで苦しんでいる人達のコメントの中でこんなコメントが

Hello friends,
Sherlock and Array hackerrank problem can be solved easily by deriving a linear equation.

引用：https://www.hackerrank.com/challenges/sherlock-and-array/forum/comments/590533

この方によると、この問題は”linear equation“、つまり線型方程式を使えばいいよということです。

線形方程式ってなんなんでしょう？
ウィキペディアによると…

線型方程式（せんけいほうていしき、linear equation）とは、線型性を持つ写像（関数・作用素）の等式で表される方程式のことである。

引用：ウィキペディア

これだけ読むとよく分からないんですが、簡単に言うと、一次方程式のことみたいです。この問題でいくと ax + by = c のような方程式をいくつか作っていけば解けてしまうよということみたいです。

さて、これを実際どう使うのか、この問題を通しての解説は、こちらの動画にまとまってました。

Repeated String HackerRank Solution [Optimal Approach]

全体の合計をsumとした時に、yを最初から順番に動かしていくと、yの左右の配列の和はsumとyを使って割り出せるので、それを使うという方法です。

条件に当てはまる時は、左右の配列の和をxを使って表すとするとそれぞれ以下のようになります。

• sum = arr.sum
• y = arr[i], 0 ≦ i < n
• sum = 2x + y

yを配列の最初から最後まで順番に動かして確認していくため、xをyの左側の要素の合計とすると、xの値は、最初は0ですが、以降yが進む毎に直前のyを足した値が、xとなっていきます。

この考え方でいくと、以下のように書けると思います。

def balancedSums(arr)
    n = arr.size
    x = 0
    sum = arr.sum
    if n == 1
        return 'YES'
    else
        arr.each do |y|
            if 2 * x == sum - y
                return 'YES'
            end
            x += y
        end
        return 'NO'
    end
end

だいぶスッキリして、無事タイムアウトにもなりませんでした。

感想

この問題はそこそこ時間を使いました。
問題を数学でいくとあれじゃんみたいな発想ができるとよりスマートにシンプルに解けるのかなあという感じです。