Subtree where leaf x is f x

from Programming (language) puzzles | Lambda the Ultimate:

You have a labeled tree, which (if my Haskell is not too rusty) is defined by:

data LTree a = Leaf String a | Branch String [LTree a]

That might correspond to a normal file system, ignoring . and .., where the leaves are files and the branches are directories. You are to produce an outline view of this tree, showing only those leaves which satisify some predicate function and the branches necessary to reach those leaves. You should assume that the tree is wide and shallow.

K2

A recursive solution using dictionaries:

st:{(./)._n,+{:[5=4:z;,/_f'[x,/:!z;y;z[]];y z;,(x;:;z);()]}[();y;x]}

A "semi-recursive" solution, also using dictionaries:

/ paths to leaves (recursive)
pl:{{:[5=4:y;,/(x,/:!y)_f'y[];,x]}[();x]}

/ subtree where leaf x is f x (non-recursive)
st:{[t;f]               / subtree of t where leaves are f
 p:pl t                 / paths to leaves
 d:t ./:p               / data at leaves
 i:&f d                 / where leaves are f
 ./[.();p i;:;d i]}     / subtree

The solution is "semi-recursive" because it packages up the recursion in the function pl, which, given a nested dictionary, returns a list of the paths to the leaves. The function st indexes out the leaf-data as a vector, applies the predicate f, and then builds a new dictionary from the paths and data which satisfy the predicate.

A non-recursive solution in "open code": t, n, and d are vector inputs representing the tree, the nodes, and the data respectively. T, N, and D are output, representing the subtree, nodes of the subtree, and data of the subtree.

t:0 0 0 2 2 2 4 4 7 7           / tree
c:#t                            / count of tree
n:!c                            / nodes (labels)
f:50<                           / predicate
l:~n _lin t                     / leaves
d:l*(c _draw 100)%l             / data (NaN at non-leaves)
b:&l&f d                        / leaves where f
p:t\'b                          / paths to leaves where f
N:{x@<x}@?,//p                  / subtree labels
z:&#N                           / zeros
i:N?/:/:p                       / path indices (renumbering)
k:,'/(-1 1_'2#,:)'i             / path indices -> pairwise links
T:@[z;*k;:;k 1]                 / subtree = renumbered tree
D:@[z%z;i[;0];:;d b]            / subtree data

K4/Q

We can do better in K4. A non-recursive solution using maps:

n:!10                         	/ nodes
t:n!0 0 0 2 2 2 4 4 7 7       	/ tree
l:~n in t                   	/ leaves
d:n!l*(10?100)%l          	/ data (NaN at non-leaves)
f:50<                       	/ predicate
N@:<N:?,/t\'&l&f d           	/ subtree nodes
T:N#t                         	/ subtree
D:N#d                         	/ subtree data

Every submap of a tree-as-map is a tree-as-map. This is not true for vectors: not every subvector of a tree-as-vector is a tree-as-vector.

Finally, K4 and Q versions of a solution where we represent a tree as a "key table": a map from a node-table to a table containing pointers-to-nodes and data-at-nodes:

K4:

t:(+(,`n)!,!10)!+(`p`d!(0 0 0 2 2 2 4 4 7 7;10?100))            / tree (as key table)
t:![t;,(in;`n;`p);0b;(,`d)!,0N]                                 / null out non-leaves
f:50<                                                           / predicate
s:?[t;,(in;`n;?,/(. t[;`p])\'?[t;,(f;`d);();`n]);0b;()]         / select subtree

Q:

t:([n:key 10];p:0 0 0 2 2 2 4 4 7 7;d:10?100)     		/ tree (as key table)
t:update d:0N from t where n in p                       	/ null out non-leaves

f:50<                                                   	/ predicate

s:select from t where
  n in distinct raze((exec p from t)scan)each 
  exec n from t where f d                               	/ select subtree

Note: scan:{x\y} added to q.k.

Thanks to Mike Rosenberg for help with the Q version, and to Felix Lungu for translation from Q to K.