Function that finds the equilibrium index of an array

equilibrium_index.rb

# === Program v1 [time complexity: O(n^2)]

class Main

  def equilibrium_index(arr)
    result = []

    arr.each_with_index do |el, i|
      j = 0
      sum_prefix = 0
      while j < i
        sum_prefix += arr[j]
        j += 1
      end

      j = i + 1
      sum_suffix = 0
      while j < arr.length
        sum_suffix += arr[j]
        j += 1
      end

      result.push(i) if sum_prefix == sum_suffix
    end

    return result
  end

end


# === Program v2 [time complexity: O(n)]

class Main

  def equilibrium_index(arr)
    result = []

    sum_suffix = 0
    arr.each { |e| sum_suffix += e }

    sum_prefix = 0
    arr.each_with_index do |el, i|
      sum_suffix -= el

      result.push(i) if sum_prefix == sum_suffix

      sum_prefix += arr[i]
    end

    return result
  end

end


# === Unit Tests

require 'minitest/autorun'
require 'minitest/reporters'
Minitest::Reporters.use! Minitest::Reporters::SpecReporter.new

class TestMain < Minitest::Test

  def setup
    @main = Main.new
    @data = [
        { arr: [-1, 3, -4, 5, 1, -6, 2, 1], idx: [1, 3, 7] }
      ]
  end

  def test_equilibrium_index
    @data.each do |el|
      assert_equal el[:idx], @main.equilibrium_index(el[:arr])
    end
  end

end