MULTI level dict: Return reference instead of the content

Question

I have a tree. I have to calculate cummulatve sum up to every leaf. And finally keep only the top N leafs by sum.

So far I got the sum and i get the leafs content back .. the problem is I want back REFERENCES to the leafs so that I delete the ones with low cumsum

Is there way to do that ?

import numpy as np

class BeamTree(dict):

  def __init__(self, *args, **kwargs):
      super(BeamTree, self).__init__(*args, **kwargs)
      self.__dict__ = self
      self.nodes = []

  def add(self, a, aa, q):
      self.nodes.append( BeamTree({ 'a': a, 'aa': aa, 'qq': q, 'qs':0 }) )
      return self.nodes[-1]

  def qsum(self, q=0):
      if len(self.nodes) == 0 : return []
      leafs = []
      for node in self.nodes:
          node['qs'] = q + node['qq']
          leafs.extend( node.qsum(node['qs']) )
          if len(node.nodes) == 0 : leafs.append(node)

      if len(leafs) > 0 : return leafs
      return []

  def generate(self, branch=3, depth=3):

      if depth < 1 : return
      for b in range(branch) :
          sym = 's' + str(np.random.randint(100))
          aix = np.random.randint(100)
          q = np.random.rand()
          node = self.add(sym, aix, q)
          node.generate(branch, depth-1)

here is a test :

In [212]: b=BeamTree(); b.generate(2,2)

  In [213]: l=b.qsum(0)

  In [214]: b
  Out[214]: 
  {'nodes': [{'a': 's80',
     'aa': 56,
     'qq': 0.673,
     'qs': 0.673,
     'nodes': [{'a': 's8', 'aa': 16, 'qq': 0.115, 'qs': 0.788, 'nodes': []}, {'a': 's64', 'aa': 10, 'qq': 0.599, 'qs': 1.272, 'nodes': []}]},
    {'a': 's67',
     'aa': 0,
     'qq': 0.900,
     'qs': 0.900,
     'nodes': [{'a': 's69', 'aa': 23, 'qq': 0.801, 'qs': 1.700, 'nodes': []}, {'a': 's8', 'aa': 41, 'qq': 0.826, 'qs': 1.726, 'nodes': []}]}]}

  In [215]: l
  Out[215]: 
  [{'a': 's8', 'aa': 16, 'qq': 0.115, 'qs': 0.788, 'nodes': []},
   {'a': 's64', 'aa': 10, 'qq': 0.599, 'qs': 1.272, 'nodes': []},
   {'a': 's69', 'aa': 23, 'qq': 0.801, 'qs': 1.700, 'nodes': []},
   {'a': 's8', 'aa': 41, 'qq': 0.826, 'qs': 1.726, 'nodes': []}]

  In [216]: del l[0]

  In [217]: l
  Out[217]: 
  [{'a': 's64', 'aa': 10, 'qq': 0.599, 'qs': 1.272, 'nodes': []},
   {'a': 's69', 'aa': 23, 'qq': 0.801, 'qs': 1.700, 'nodes': []},
   {'a': 's8', 'aa': 41, 'qq': 0.826, 'qs': 1.726, 'nodes': []}]

  In [218]: b
  Out[218]: 
  {'nodes': [{'a': 's80',
     'aa': 56,
     'qq': 0.673,
     'qs': 0.673,
     'nodes': [{'a': 's8', 'aa': 16, 'qq': 0.115, 'qs': 0.788, 'nodes': []}, {'a': 's64', 'aa': 10, 'qq': 0.599, 'qs': 1.272, 'nodes': []}]},
    {'a': 's67',
     'aa': 0,
     'qq': 0.900,
     'qs': 0.900,
     'nodes': [{'a': 's69', 'aa': 23, 'qq': 0.801, 'qs': 1.700, 'nodes': []}, {'a': 's8', 'aa': 41, 'qq': 0.826, 'qs': 1.726, 'nodes': []}]}]}

  In [219]: 

Answer 1

I want back REFERENCES to the leafs so that I delete the ones with low cumsum

Actually, you got references to the leaves, as in Python "everything" is a reference, but what you really need is the reference to the parent's nodes list, so that you can remove the leaf reference from that list.

So something like this:

if len(node.nodes) == 0: leafs.append((self.nodes, node))

Now you have (list, node) tuples, and to remove such leaf, you would do something like this:

for nodes, leaf in leafs:
    if <your condition for removal>:
        nodes.remove(leaf)