Skip to content

Instantly share code, notes, and snippets.

@lshifr
Created October 6, 2015 22:22
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save lshifr/8f0da2ceec5fe8a109c7 to your computer and use it in GitHub Desktop.
Save lshifr/8f0da2ceec5fe8a109c7 to your computer and use it in GitHub Desktop.
Creating nested structure from flat, for Mathematica expressions
BeginPackage["Unflatten`"]
Unflatten::usage = "Unflatten[h[elems], posints, heads] creates a nested structure by wrapping
the elements in positions intervals posints in heads heads";
UnflattenNested::usage = "UnflattenNested[expr, {startposlist, endposlist}, head] wraps elements
of expr at positions starting at startposlist_i and ending at endposlist_i in head head";
Begin["`Private`"]
ClearAll[listSplit,reconstructIntervals,groupElements,groupPositions,processPosList,groupElementsNested];
listSplit[x_List, lengthlist_List,headlist_List]:=
MapThread[
#1@@Take[x,#2]&,
{headlist,Transpose[{Most[#]+1,Rest[#]}&[FoldList[Plus,0,lengthlist]]]}
];
reconstructIntervals[listlen_Integer,ints_List]:=
Module[{missed,startint,lastint},
startint = If[ints[[1,1]]==1,{},{1,ints[[1,1]]-1}];
lastint = If[ints[[-1,-1]]==listlen, {},{ints[[-1,-1]]+1,listlen}];
missed =
Map[
If[#[[2,1]]-#[[1,2]]>1,{#[[1,2]]+1,#[[2,1]]-1},{}]&,
Partition[ints,2,1]
];
missed = Join[missed,{lastint}];
Prepend[Flatten[Transpose[{ints,missed}],1],startint]
];
groupElements[
lst_List,poslist_List,headlist_List
] /; And[OrderedQ[Flatten[Sort[poslist]]],Length[headlist]==Length[poslist]]:=
Module[{totalheadlist,allints,llist},
totalheadlist =
Append[
Flatten[Transpose[{Array[Sequence&,{Length[headlist]}],headlist}],1],
Sequence
];
allints = reconstructIntervals[Length[lst],poslist];
llist = Map[If[#==={},0,1-Subtract@@#]&,allints];
listSplit[lst, llist,totalheadlist]
];
(* To work on general heads, we need this *)
groupElements[h_[x__],poslist_List,headlist_List]:=
h[Sequence@@groupElements[{x},poslist,headlist]];
(* If we have a single head *)
groupElements[expr_,poslist_List,head_]:=
groupElements[expr,poslist,Table[head,{Length[poslist]}]];
groupPositions[plist_List]:=
Reap[Sow[Last[#],{Most[#]}]&/@plist,_,List][[2]];
processPosList::unmtch = "The starting and ending position lists `1` and `2` don't match";
processPosList[{openlist_List,closelist_List}]:=
Module[{opengroup,closegroup,poslist},
{opengroup,closegroup} = groupPositions/@{openlist,closelist};
poslist = Transpose[Transpose[Sort[#]]&/@{opengroup,closegroup}];
If[UnsameQ @@ poslist[[1]],
Message[processPosList::unmtch,openlist,closelist];
{},
(* else *)
Transpose[{poslist[[1,1]],Transpose/@ Transpose[poslist[[2]]]}]
]
];
groupElementsNested[nested_,{openposlist_List,closeposlist_List},head_]/;Head[head]=!=List:=
Fold[
Function[{x,y},MapAt[groupElements[#,y[[2]],head]&,x,{y[[1]]}]],
nested,
Sort[processPosList[{openposlist,closeposlist}],Length[#2[[1]]] < Length[#1[[1]]]&]
];
Unflatten[args___]:= groupElements[args];
UnflattenNested[args___]:=groupElementsNested[args]
End[]
EndPackage[]
@valerio-marra
Copy link

valerio-marra commented Nov 3, 2017

Hi, is this equivalent to the new ArrayReshape?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment