CCC 2018 S4 - Balanced Trees

1# By Daniel Zhang, Pinetree Secondary School2 3import math4 5n = int(input())6 7 8memo = {}9#note that this is O(W) = O(10^9)10#we can optimize by noticing that the same weight subtrees will get called many times11#take recurse(4) for example12#it can have k = 2, 3, 4 subtres13 14#2 Subtrees: -> each subtree has weight 4//2 = 2 15#3 subtrees: -> each subtree has weight 4//3 = 116#4 subtrees: -> each subtree has weight 4//4 = 117#as we can see, the number of PBT's of weight 4 is F(2) + F(1) + F(1)18#the F(1) is called twice, we can isntead multiply F(1)19#now it comes down to finding the number of F(i) for each i20 21#for a value v, the range of i which satisfies it is n / v+1 <= i < n / v22#after we get the count, we just need to find each valid value23#to do this we can 'jump' to the next range24 25#wuganggame's code (finally makes sense!)26#d is the value of each subtree27#r is the remaining, and so r/d is additional i-values that we have where W/i is still d28#then we add one similar to next_i29#   int k = 2;30#   while (k <= n) {31#     int d = n / k;32#     int r = n % k;33#     int count = r / d + 1;34#     ans += count * f(d);35#     k += count;36#   }37def recurse(w): #number of subtrees for a pbt of  weight w38    if w == 1:39        return 140    if w in memo:41        return memo[w]42 43    out = 044 45    i = 246    while i <= w:47        v = w // i48        numV = w // v - w // (v+1)49        out += numV * recurse(v)50        51        #this is the key part, and also hardest to understand52        #some v values won't show up53        #for example in F(10) there is no 10//i = 454        #so we should skip that and go from 5 -> 355        #the current weight of each subtree is w // i56        #so the number of times that V can fit into W is the largest i that produces this V57        #example: W = 10, i = 458        #10//4 = 2 = V59        # 10 // 2 is 5, so 5 is the maximum i for which W//i equals V60        #thus, i+1 will result in a new V61        #another example: W = 12, i = 362        #V = 12//3 = 463        #w // v = 3, and so next_i = 3+1 = 464        #so indeed 12//4 gives a new V value of 3, from the previous V value of 465        next_i = w // v + 166 67        i = next_i68 69    memo[w] = out70    return out71 72print(recurse(n))