Generalizing lua functions for n-dimensional arrays

Question

The following MWE initializes, increments and prints 2-D and 3-D numeric arrays. It could easily be extended to handle 4-D etc arrays but it would be more elegant to have generic functions that took the number of dimensions (up to say 7) as a parameter or which determined the dimension from the number of input parameters. What would be a good way to implement that?

MWE:

local i_low, i_high = 2, 4
local i_range = i_high - i_low - 1

local j_low, j_high = 3, 7
local j_range = j_high - j_low - 1

local k_low, k_high = 1, 3
local k_range = k_high - k_low - 1

local myArray_two = {}
local myArray_three = {}

function initArray_two(t, i_low, i_high, j_low, j_high, value)
    local i, j = 0, 0

    for i = i_low, i_high, 1 do
        for j = j_low, j_high, 1 do
            local idx = j * i_range + i
            t[idx] = value
        end
    end
end

function initArray_three(t, i_low, i_high, j_low, j_high, k_low, k_high, value)
    local i, j, k = 0, 0, 0

    for i = i_low, i_high, 1 do
        for j = j_low, j_high, 1 do
            for k = k_low, k_high, 1 do
                local idx = k*j_range*i_range + j*i_range + i
                t[idx] = value
            end
        end
    end
end

function incrValue_two(t, i, j, value)
    assert(i>=i_low and i <= i_high)
    assert(j>=j_low and j <= j_high)

    local idx = j * i_range + i
    t[idx] = t[idx] + value    
end

function incrValue_three(t, i, j, k, value)
    assert(i>=i_low and i <= i_high)
    assert(j>=j_low and j <= j_high)
    assert(k>=k_low and k <= k_high)

    local idx = k*j_range*i_range + j*i_range + i
    t[idx] = t[idx] + value
end

function printArray_two(t, title, i_low, i_high, j_low, j_high)
    local i, j = 0, 0
    print(title.."\n")

    for i = i_low, i_high, 1 do 
        for j = j_low, j_high, 1 do                    
            local idx = j * i_range + i     
            print(i.."\t"..j.."\t"..t[idx])    
        end      
        print("\n")    
    end 

end  

function printArray_three(t, title, i_low, i_high, j_low, j_high, k_low, k_high)  
    local i, j, k = 0, 0, 0     
    print(title.."\n")

    for i = i_low, i_high, 1 do 
        for j = j_low, j_high, 1 do        
            for k = k_low, k_high, 1 do              
                local idx = k*j_range*i_range + j*i_range + i       
                print(i.."\t"..j.."\t"..k.."\t"..t[idx].."\n")
            end     
        end      
    end 

end  

initArray_two(myArray_two, i_low, i_high, j_low, j_high, 0)
initArray_three(myArray_three, i_low, i_high, j_low, j_high, k_low, k_high, 1)

incrValue_two(myArray_two, 2, 3, 11)
incrValue_two(myArray_two, 2, 3, 13)
incrValue_two(myArray_two, 4, 7, 5)
printArray_two(myArray_two, "A 2-D Array", i_low, i_high, j_low, j_high)

incrValue_three(myArray_three, 2, 3, 1, 9)
incrValue_three(myArray_three, 2, 3, 1, 17)
printArray_three(myArray_three, "A 3-D Array", i_low, i_high, j_low, j_high, k_low, k_high)

Answer 1

You can generalize the functions using ... in the function definition. This will capture your repeating arguments, ideally your ranges.

function initArray(t, default_value, range, ...)
  local args = {...} -- our next ranges
  if args[1] then    -- if we have more ranges recurse
    for i = range[1], range[2], 1 do
      t[i] = initArray({}, default_value, table.unpack(args))
    end
  else               -- if we dont have more ranges set default values
    for i = range[1], range[2], 1 do
      t[i] = default_value
    end
  end
  return t
end

I also change your calls to create pairs of high and low range values

initArray(myArray_two, 0, {i_low, i_high}, {j_low, j_high})

Here is the whole code:

local i_low, i_high = 2, 4
local j_low, j_high = 3, 7
local k_low, k_high = 1, 3

local myArray_two = {}
local myArray_three = {}


function initArray(t, default_value, range, ...)
  local args = {...}
  if args[1] then
    for i = range[1], range[2], 1 do
      t[i] = initArray({}, default_value, table.unpack(args))
    end
  else
    for i = range[1], range[2], 1 do
      t[i] = default_value
    end
  end
  return t
end


function incrValue(t, value, ...)
  local args = {...}
  assert(#args >= 1)

  local el = t
  for _,v in ipairs(args) do
    if type(el[v]) == 'table' then
      el = el[v]
    else
      el[v] = el[v] + value 
    end
  end
end

function printArray(t, title)
  print(title .. "\n")

  for k, v in pairs(t) do     
    if type(v) == 'table' then
      recurPrintArray(k , v)
    else
      print(k .. "\t" .. v)
    end
    print("\n")    
  end 
end  

function recurPrintArray(s, t)
  for k, v in pairs(t) do     
    if type(v) == 'table' then
      recurPrintArray(s .. "\t" .. k, v)
    else
      print(s .. "\t" .. k .. "\t" .. v)
    end
  end 

end

initArray(myArray_two, 0, {i_low, i_high}, {j_low, j_high})
initArray(myArray_three, 1, {i_low, i_high}, {j_low, j_high}, {k_low, k_high})

incrValue(myArray_two, 11, 2, 3)
incrValue(myArray_two, 13, 2, 3)
incrValue(myArray_two, 5, 4, 7)
printArray(myArray_two, "A 2-D Array", i_low, i_high, j_low, j_high)

incrValue(myArray_three, 9, 2, 3, 1)
incrValue(myArray_three, 17, 2, 3, 1)
printArray(myArray_three, "A 3-D Array", i_low, i_high, j_low, j_high, k_low, k_high)