Map a function at a given level of a tree

I feel a bit dumb asking this very basic question, but I have looked and looked but cannot figure it out.

Imagine I have a “tree”, in the sense of a container of containers of containers… Such as

julia> tree = [[[1,2],[3,4]],[[5,6,7],[8]],[[9,10]]]
3-element Array{Array{Array{Int64,1},1},1}:
 [[1, 2], [3, 4]]
 [[5, 6, 7], [8]]
 [[9, 10]]

And say I want to map a function f at a given level of this tree. For the first level I just do map(f, tree). For the second level I can do map.(f, tree). For the third level I run out of ideas and I am forced to spell it out with a loop. Is there something like a mapat(f, tree, level)?

Thanks for any pointers!

It can be done with a recursive definition like this:

2 Likes

Thanks @tkf, but I don’t think that is quite what I need. rebroadcast applies f at every level down to n (right?). I want to apply it strictly at level n. But your idea is great. I can do something similar, like

julia> mapat(f, t, level) = level == 0 ? f(t) : mapat.(f, t, level-1);

julia> mapat(string, tree, 3)
3-element Array{Array{Array{String,1},1},1}:
 [["1", "2"], ["3", "4"]]
 [["5", "6", "7"], ["8"]]
 [["9", "10"]]

Thanks!

EDIT: I had misunderstood! No, your rebroadcast is just like the above mapat, only written in a different form. Thanks again.

1 Like

The solution is this I think

function MagicMapAtLevel(level,func,tree)
    if level == 1
        return map(func,tree)
    else
        for subtree in tree
          MagicMapAtLevel(level-1,func,subtree)
        end
    end
end

tree = [[[1,2],[3,4]],[[5,6,7],[8]],[[9,10]]]

println("level 1")
MagicMapAtLevel(1,println,tree)

println("level 2")
MagicMapAtLevel(2,println,tree)

println("level 3")
MagicMapAtLevel(3,println,tree)

with output

level 1
Array{Int64,1}[[1, 2], [3, 4]]
Array{Int64,1}[[5, 6, 7], [8]]
Array{Int64,1}[[9, 10]]
level 2
[1, 2]
[3, 4]
[5, 6, 7]
[8]
[9, 10]
level 3
1
2
3
4
5
6
7
8
9
10

If you don’t need the level information you can omit level as parameter. For instance when using arrays with Int64 as base type the function treeMap could have the following form:

function treeMap(f, tree)
   if typeof(tree) == Array{Int64,1}
      return map(f, tree)
   else
      return map( t -> treeMap(f,t), tree)
   end

end

So for instance we get for

treeMap(x -> x^2, tree)

the following result

3-element Array{Array{Array{Int64,1},1},1}:
 [[1, 4], [9, 16]]
 [[25, 36, 49], [64]]
 [[81, 100]]

The requirement of the OP is to specify a level. For more general solution, maybe I’d go with something like

julia> mapover(f, iselement, x) =
           iselement(x) ? f(x) : map(e -> mapover(f, iselement, e), x)
mapover (generic function with 1 method)

julia> mapover(string, x -> x isa Number, [[1, [2]], (3, (four = 4,))])
2-element Array{Any,1}:
 Any["1", ["2"]]
 ("3", (four = "4",))
1 Like