code: def get_topic_imbalance_stats(brokers, topics):
      
      extra_partition_cnt_per_broker = defaultdict(int)
      tot_brokers = len(brokers)
      # Sort the brokers so that the iteration order is deterministic.
      sorted_brokers = sorted(brokers, key=lambda b: b.id)
      for topic in topics:
          # Optimal partition-count per topic per broker
          total_partition_replicas = \
              len(topic.partitions) * topic.replication_factor
          opt_partition_cnt, extra_partitions_allowed = \
              compute_optimum(tot_brokers, total_partition_replicas)
          # Get extra-partition count per broker for each topic
          for broker in sorted_brokers:
              partition_cnt_broker = broker.count_partitions(topic)
              extra_partitions, extra_partitions_allowed = \
                  get_extra_element_count(
                      partition_cnt_broker,
                      opt_partition_cnt,
                      extra_partitions_allowed,
                  )
              extra_partition_cnt_per_broker[broker.id] += extra_partitions
  
      # Net extra partitions over all brokers
      net_imbalance = sum(six.itervalues(extra_partition_cnt_per_broker))
      return net_imbalance, extra_partition_cnt_per_broker
docs: Return count of topics and partitions on each broker having multiple
      partitions of same topic.
  
      :rtype dict(broker_id: same-topic-partition count)
      Example:
      Total-brokers (b1, b2): 2
      Total-partitions of topic t1: 5
      (b1 has 4 partitions), (b2 has 1 partition)
      opt-count: 5/2 = 2
      extra-count: 5%2 = 1
      i.e. 1 broker can have 2 + 1 = 3 partitions
      and rest of brokers can have 2 partitions for given topic
      Extra-partition or imbalance:
      b1: current-partitions - optimal-count = 4 - 2 - 1(extra allowed) = 1
      Net-imbalance = 1