Perry the Panda is eating some bamboo over a five-day period from Monday to Friday (inclusive). On Monday, he eats
14
pieces of bamboo. Each following day, Perry eats either one less than three times the previous day or one more than the previous day, with equal probability. Compute the expected number of pieces of bamboo Perry has eaten throughout the week after the end of Friday. (AOPS link)
This problem can be solved in different ways and it serves an excellent example to practise tree manipulation in Wolfram language.

Analysis

Direct computation

Regardless of the probability, there are two choices for the number of bamboos to consume on the next day. NestTree is the suitable function to enumerate all states:
In[]:=
dietTree=NestTree[{#+1,#*3-1}&,14,4]
The corresponding probability tree is also a complete binary tree:
In[]:=
dietTreeProb=NestTree[{#*1/2,#*1/2}&,1,4]
By definition the expected value of discrete events is simply
, thus the result is the sum of node-wise product between the two trees. Usually you cannot multiply two trees node-wisely without checking if the trees are identical in terms of structure. In our case we can omit the process because the diet tree and its probability tree are generated with same NestTree structure and we can implement the tree-like dot product with following code:
In[]:=
myTreeMultiply[t1_,t2_]:=TreeMap[(#1*​​TreeExtract[t2,{#2},TreeData][[1]])&,t1,All->{"Data","Position"}]
In[]:=
myTreeMultiply[dietTree,dietTreeProb]
TreeFold gives a one-liner to compute the sum of all elements in the tree above:
In[]:=
TreeFold[Total[Flatten[{##}]]&,%]//Framed
Out[]=
434
434 is the expected value that we are looking for.

Better solution for contestants

AIME problems usually requires contestant to solve the problems in few minutes. It is not advised to sum up all element blindly in random order. Lets look at the tree or its NestTree argument again:
In[]:=
dietTree
Out[]=
Notice that {x + 1, x*3 - 1} indicates the sum of two direct children is 4 times the value of parent node. Therefore if we compute the node sum in Breath-first-search (BFS) order, we obtain a geometric sequence:
In[]:=
res1=Reap[TreeScan[Sow,dietTree,​​TreeTraversalOrder->{"BreadthFirst","PreOrder","LeftRight"}]][[2,1]]
Out[]=
{14,15,41,16,44,42,122,17,47,45,131,43,125,123,365,18,50,48,140,46,134,132,392,44,128,126,374,124,368,366,1094}
In[]:=
hSum1=Table[Total[res1[[2^k;;2^(k+1)-1]]],{k,0,Log2[Length[res1]-1]}]
Out[]=
{14,56,224,896,3584}
In[]:=
(#2/#1&)@@@Partition[hSum1,2,1]
Out[]=
{4,4,4,4}
Therefore the final result can written as the following assume we still have the same probability tree:
In[]:=
(14*4^#).(1/2^#)&[Range[0,4]]//Framed
Out[]=
434
For instance, we can extrapolate the solution to an expected value after a 10-day period:
In[]:=
(14*4^#).(1/2^#)&[Range[0,10]]
Out[]=
28658
Or to a closed form solution:
In[]:=
res=Sum[(14*4^k)*(1/2^k),{k,0,n}]
Out[]=
14(-1+
1+n
2
)
In[]:=
res/.{n->{4,10}}
Out[]=
{434,28658}

Probability as a variable

The above BFS works well when the row (items with same tree level) sums in the diet tree forms a geometric sequence and each probability in each row of the probability tree is identical. Lets generalize the probability tree by replacing
1/2
with a variable
p
,
0<=p<=1
.
In[]:=
dietTreeProbF[p_]:=NestTree[{#*p,#*(1-p)}&,1,4]
Span the probability of the occurrence of each event on the binary tree:
In[]:=
TreeMap[Expand,dietTreeProbF[p],TreeLayout->Left]
Use the definition of expected value we can see the sum of elements in the same tree level does not cancel
p
. So we opt to direct sum method:
In[]:=
Tree[res2=myTreeMultiply[dietTree,dietTreeProbF[p]],TreeLayout->Left]
The expectation of number of bamboo consumed over a 5-day period:
Visualize the dependence of the expected value upon probability of next-day events:

Summary

◼
  • NestTree is useful to construct event trees with dependence upon previous outcome.
  • ◼
  • TreeScan with options can control the order for visiting node; useful to implement BFS or DFS.
  • ◼
  • Use TreeMap to construct a Dot-product-like operation between trees with identical level structure.
  • ◼
  • Acknowledge to my collogue Ian Ford for valuable suggestion on tree traversal.
  • CITE THIS NOTEBOOK

    Solve an OTIS AIME 24' problem with tree traversal​
    by Shenghui Yang​
    Wolfram Community, STAFF PICKS, March 3, 2024
    ​https://community.wolfram.com/groups/-/m/t/3133801