Compressed Row Storage, CRS for sparse matrix

crs_test.rb

require 'pp'

class CRS
  attr_reader :col_indexes, :row_pointers
  def initialize data
    # @vals = [] not used
    @col_indexes = []
    @row_pointers = []
    @n = nil
    counter = 0
    data.each do |i|
      i.each do |j|
        @col_indexes.push j
        @row_pointers[counter] = @col_indexes.size - 1 if @row_pointers[counter].nil?
      end  
      counter = counter + 1
    end
    row_pointers.push(counter + 1)
    @n = counter
  end

  def trans_square
    results = []
    @n.times do |i|
      result = {}
      s = @row_pointers[i]
      e = @row_pointers[i+1]
      cols = @col_indexes[s..(e-1)]
      @n.times do |j|
        val = cols.inject(0) do |sum,k|
          if self.elm(j, k) != 0
            sum + 1
          else
            sum
          end
        end
        result[j] = val if val != 0 
      end
      results.push result
    end
    results
  end
  
  def square
    results = []
    @n.times do |i|
      result = {}
      s = @row_pointers[i]
      e = @row_pointers[i+1]
      cols = @col_indexes[s..(e-1)]
      @n.times do |j|
        val = cols.inject(0) do |sum,k|
          if self.elm(k, j) != 0
            sum + 1
          else
            sum
          end
        end
        result[j] = val if val != 0 
      end
      results.push result
    end
    results
  end

  def elm(x,y)
    s = @row_pointers[x]
    e = @row_pointers[x+1]
    if @col_indexes[s..(e-1)].include? y
      1
    else
      0
    end
  end
end
data = [[1,2],[0],[0]]
a = CRS.new(data)
pp a
pp a.elm(1,1)
pp a.elm(0,2)
pp a.square

data = [[1,2],[0]]
a = CRS.new(data)
pp a
pp a.trans_square