q2d_2indent.py

class Constraint:
  """Class representing a tree constraint such as those used by the BUILD algorithm.
  Stored constraint becomes (a, b) < (c, d) where a, b, c, d are the four arguments provided.
  Access Constraint.text to get the constraint as a human-readable string.
  """

  def __init__(self, a: str, b: str, c: str, d:str):
    self.items = [a, b, c, d]
    self.text = f'({a}, {b}) < ({c}, {d})'


class TreeNode:
  """Class representing a node in a tree. Takes two arguments:
  - children: A list of TreeNodes representing children (provide [] for a leaf)
  - name: The node's name (required for a leaf, optional otherwise)
  """

  def __init__(self, children: list, name: str = ''):
    self.children = children
    self.name = name
    self.isLeaf = children == []
  
  def traverse(self) -> (list, list, list):
    """Returns three values in a tuple:
    - A list of constraints representing the subtree below the current node
    - A 'flattened' list of the named nodes beneath the current node in the tree
      ie, a node with children (a node connected to (a, b), c) would return [a, b, c]
    - A list of pairs of nodes that would need to be linked together by constraints created
      by any parent of the current node
    """

    # Leaves return just themselves
    if (self.isLeaf):
      return ([], [self.name], [])
    
    numChildren = len(self.children)


    # Traverse children
    leaves = []
    internals = []

    childFlatLists = []
    constraintLists = []

    for i in self.children:
      (constraints, flatList, pairsToConnect) = i.traverse()

      if (len(flatList) == 1):
        leaves.append(flatList[0])
      
      else:
        internals.append((constraints, flatList, pairsToConnect))
        childFlatLists.append(flatList)
        constraintLists.append(constraints)

    numLeaves = len(leaves)
    numInternals = len(internals)


    # Generate constraints from data returned by children
    myConstraints = []
    myFlatList = leaves
    myPairsToConnect = []
    childPairsToConnect = []

    # If we're only connected to leaves, just return the flattened list of leaves,
    #   making sure the parent connects them
    if (numInternals == 0):
      return ([], leaves, [(leaves[i], leaves[i+1]) for i in range(numLeaves-1)])

    # If there are multiple internal children, link them all with constraints
    elif (numInternals > 1):
      # As we are connecting two internal children at a time, asking the parent to make one constraint for
      #   each would result in an unnecessary 'loop' - therefore we can omit the last one
      for i in range(numInternals-1):
        
        # Ensure that the current internal is linked to the next one in any parent of the current node
        thisFlatList = internals[i][1]
        nextFlatList = internals[i+1][1]
        myPairsToConnect.append((thisFlatList[0], nextFlatList[0]))

        # Add any pairsToConnect from the current internal to a central list
        childPairsToConnect.extend(internals[i][2])
    
    # Add pairsToConnect for any internal children not covered above
    #   (either a single internal child skipped by the if, or the final one which was not covered by the loop)
    childPairsToConnect.extend(internals[numInternals-1][2])
    
    # Link any leaves to an internal
    # (If we've got this far without returning, we have at least one internal child)
    firstFlatList = internals[0][1]
    for i in leaves:
      # If we have pairsToConnect from internal children, use those
      if (len(childPairsToConnect) > 0):
        connectingPair = childPairsToConnect.pop(0)
      
      else:
        # Otherwise, just use the first two leaves in the first internal
        connectingPair = firstFlatList
      
      # Generate constraint, adding the right hand side to the list of pairs to connect for any parent
      myConstraints.append(Constraint(connectingPair[0], connectingPair[1], connectingPair[0], i))
      myPairsToConnect.append((connectingPair[0], i))
    
    # If there are any more pairsToConnect from children
    while (len(childPairsToConnect) > 0):
      toConnect = childPairsToConnect.pop(0)

      # If we have any leaves, just connect to the first one
      if (numLeaves != 0):
        toConnectTo = leaves[0]
      
      else:
        # Otherwise, find a pair from any internal that doesn't include the pair we're connecting
        # There will always be at least two children in a valid tree, so
        #   if there are no leaves there will be enough internals
        toConnectTo = next(
          flatList[0]
          for flatList in childFlatLists
          if ((toConnect[0] not in flatList) and (toConnect[1] not in flatList))
        )
      
      # Generate the constraint
      #   (no need to add to myPairsToConnect as the right hand side is superfluous)
      myConstraints.append(Constraint(toConnect[0], toConnect[1], toConnect[0], toConnectTo))
    
    # Prepare output & return
    for i in childFlatLists:
      myFlatList.extend(i)
    
    for i in constraintLists:
      myConstraints.extend(i)

    return (myConstraints, myFlatList, myPairsToConnect)
  
  def getConstraints(self):
    """Returns a list of constraints representing the (sub)tree under the current node
    (Executes traverse() and returns only the constraints)
    """

    return self.traverse()[0]