text: # -*- coding: utf-8 -*-
  #
  # lukes.py - Lukes tree partitioning algorithm
  # Author: Federico Rosato <federico.rosato@supsi.ch>
  """Lukes Algorithm for exact optimal weighted tree partitioning."""
  
  from copy import deepcopy
  from functools import lru_cache
  from random import choice
  
  import networkx as nx
  from networkx.utils import not_implemented_for
  
  __all__ = ['lukes_partitioning']
  
  D_EDGE_W = 'weight'
  D_EDGE_VALUE = 1.0
  D_NODE_W = 'weight'
  D_NODE_VALUE = 1
  PKEY = 'partitions'
  CLUSTER_EVAL_CACHE_SIZE = 2048
  
  
  def _split_n_from(n: int, min_size_of_first_part: int):
      # splits j in two parts of which the first is at least
      # the second argument
      assert n >= min_size_of_first_part
      for p1 in range(min_size_of_first_part, n + 1):
          yield p1, n - p1
  
  
  def lukes_partitioning(G,
                         max_size: int,
                         node_weight=None,
                         edge_weight=None) -> list:
  
      """Optimal partitioning of a weighted tree using the Lukes algorithm.
  
      This algorithm partitions a connected, acyclic graph featuring integer
      node weights and float edge weights. The resulting clusters are such
      that the total weight of the nodes in each cluster does not exceed
      max_size and that the weight of the edges that are cut by the partition
      is minimum. The algorithm is based on LUKES[1].
  
      Parameters
      ----------
      G : graph
  
      max_size : int
          Maximum weight a partition can have in terms of sum of
          node_weight for all nodes in the partition
  
      edge_weight : key
          Edge data key to use as weight. If None, the weights are all
          set to one.
  
      node_weight : key
          Node data key to use as weight. If None, the weights are all
          set to one. The data must be int.
  
      Returns
      -------
      partition : list
          A list of sets of nodes representing the clusters of the
          partition.
  
      Raises
      -------
      NotATree
          If G is not a tree.
      TypeError
          If any of the values of node_weight is not int.
  
      References
      ----------
      .. Lukes, J. A. (1974).
         "Efficient Algorithm for the Partitioning of Trees."
         IBM Journal of Research and Development, 18(3), 217–224.
  
      """
      # First sanity check and tree preparation
      if not nx.is_tree(G):
          raise nx.NotATree('lukes_partitioning works only on trees')
      else:
          if nx.is_directed(G):
              root = [n for n, d in G.in_degree() if d == 0]
              assert len(root) == 1
              root = root[0]
              t_G = deepcopy(G)
          else:
              root = choice(list(G.nodes))
              # this has the desirable side effect of not inheriting attributes
              t_G = nx.dfs_tree(G, root)
  
      # Since we do not want to screw up the original graph,
      # if we have a blank attribute, we make a deepcopy
      if edge_weight is None or node_weight is None:
          safe_G = deepcopy(G)
          if edge_weight is None:
              nx.set_edge_attributes(safe_G, D_EDGE_VALUE, D_EDGE_W)
              edge_weight = D_EDGE_W
          if node_weight is None:
              nx.set_node_attributes(safe_G, D_NODE_VALUE, D_NODE_W)
              node_weight = D_NODE_W
      else:
          safe_G = G
  
      # Second sanity check
      # The values of node_weight MUST BE int.
      # I cannot see any room for duck typing without incurring serious
      # danger of subtle bugs.
      all_n_attr = nx.get_node_attributes(safe_G, node_weight).values()
      for x in all_n_attr:
          if not isinstance(x, int):
              raise TypeError('lukes_partitioning needs integer '
                              'values for node_weight ({})'
                              .format(node_weight))
  
      # SUBROUTINES -----------------------
      # these functions are defined here for two reasons:
      # - brevity: we can leverage global "safe_G"
      # - caching: signatures are hashable
  
      @not_implemented_for('undirected')
      # this is intended to be called only on t_G
      def _leaves(gr):
          for x in gr.nodes:
              if not nx.descendants(gr, x):
                  yield x
  
      @not_implemented_for('undirected')
      def _a_parent_of_leaves_only(gr):
          tleaves = set(_leaves(gr))
          for n in set(gr.nodes) - tleaves:
              if all([x in tleaves for x in nx.descendants(gr, n)]):
                  return n
  
      @lru_cache(CLUSTER_EVAL_CACHE_SIZE)
      def _value_of_cluster(cluster: frozenset):
          valid_edges = [e for e in safe_G.edges
                         if e[0] in cluster and e[1] in cluster]
          return sum([safe_G.edges[e][edge_weight] for e in valid_edges])
  
      def _value_of_partition(partition: list):
          return sum([_value_of_cluster(frozenset(c)) for c in partition])
  
      @lru_cache(CLUSTER_EVAL_CACHE_SIZE)
      def _weight_of_cluster(cluster: frozenset):
          return sum([safe_G.nodes[n][node_weight] for n in cluster])
  
      def _pivot(partition: list, node):
          ccx = [c for c in partition if node in c]
          assert len(ccx) == 1
          return ccx[0]
  
      def _concatenate_or_merge(partition_1: list, partition_2: list,
                                x, i, ref_weigth):
  
          ccx = _pivot(partition_1, x)
          cci = _pivot(partition_2, i)
          merged_xi = ccx.union(cci)
  
          # We first check if we can do the merge.
          # If so, we do the actual calculations, otherwise we concatenate
          if _weight_of_cluster(frozenset(merged_xi)) <= ref_weigth:
              cp1 = list(filter(lambda x: x != ccx, partition_1))
              cp2 = list(filter(lambda x: x != cci, partition_2))
  
              option_2 = [merged_xi] + cp1 + cp2
              return option_2, _value_of_partition(option_2)
          else:
              option_1 = partition_1 + partition_2
              return option_1, _value_of_partition(option_1)
  
      # INITIALIZATION -----------------------
      leaves = set(_leaves(t_G))
      for lv in leaves:
          t_G.nodes[lv][PKEY] = dict()
          slot = safe_G.nodes[lv][node_weight]
          t_G.nodes[lv][PKEY][slot] = [{lv}]
          t_G.nodes[lv][PKEY][0] = [{lv}]
  
      for inner in [x for x in t_G.nodes if x not in leaves]:
          t_G.nodes[inner][PKEY] = dict()
          slot = safe_G.nodes[inner][node_weight]
          t_G.nodes[inner][PKEY][slot] = [{inner}]
  
      # CORE ALGORITHM -----------------------
      while True:
          x_node = _a_parent_of_leaves_only(t_G)
          weight_of_x = safe_G.nodes[x_node][node_weight]
          best_value = 0
          best_partition = None
          bp_buffer = dict()
          x_descendants = nx.descendants(t_G, x_node)
          for i_node in x_descendants:
              for j in range(weight_of_x, max_size + 1):
                  for a, b in _split_n_from(j, weight_of_x):
                      if a not in t_G.nodes[x_node][PKEY].keys() \
                              or b not in t_G.nodes[i_node][PKEY].keys():
                          # it's not possible to form this particular weight sum
                          continue
  
                      part1 = t_G.nodes[x_node][PKEY][a]
                      part2 = t_G.nodes[i_node][PKEY][b]
                      part, value = _concatenate_or_merge(part1, part2,
                                                          x_node, i_node, j)
  
                      if j not in bp_buffer.keys() or bp_buffer[j][1] < value:
                          # we annotate in the buffer the best partition for j
                          bp_buffer[j] = part, value
  
                      # we also keep track of the overall best partition
                      if best_value <= value:
                          best_value = value
                          best_partition = part
  
              # as illustrated in Lukes, once we finished a child, we can
              # discharge the partitions we found into the graph
              # (the key phrase is make all x == x')
              # so that they are used by the subsequent children
              for w, (best_part_for_vl, vl) in bp_buffer.items():
                  t_G.nodes[x_node][PKEY][w] = best_part_for_vl
              bp_buffer.clear()
  
          # the absolute best partition for this node
          # across all weights has to be stored at 0
          t_G.nodes[x_node][PKEY][0] = best_partition
          t_G.remove_nodes_from(x_descendants)
  
          if x_node == root:
              # the 0-labeled partition of root
              # is the optimal one for the whole tree
              return t_G.nodes[root][PKEY][0]
  