code: def rebalance_distributions(
      dist1_counts: np.array, dist2_counts: np.array
  ):
      """Rebalance the distributions as if dist1 was a even distribution.
  
      This is a util function for an unbalanced version categorical drift scoring methods. The rebalancing is done
      so in practice all categories of the distributions are treated with the same weight.
  
      The function redefines the dist1_counts to have equal counts for all categories, and then redefines the
      dist2_counts to have the same "change" it had from dist1_counts, but relative to the new dist1_counts. In
      addition, we add 1 to all counts. This is inspired from the properties of Beta distribution for a bernoulli
      distribution parameter estimation, see
      http://www.ece.virginia.edu/~ffh8x/docs/teaching/esl/2020-04/farnoud-slgm-chap03.pdf for additional details.
  
      Example:
          if dist1_counts was [9000, 1000] and dist2_counts was [8000, 2000]. This means that if we treat dist1 as a
          even distribution, the new counts should be dist1_counts = [5000, 5000].
          The relative change of dist2 from dist1 was a decrease of ~11% in the first category and an increase of
          ~200% in the second category. The new dist2_counts should be [4445, 9995].
          # When re-adjusting to the original total num_samples of dist2, the new dist2_counts should be [3078, 6922]
      """
      # downsizing to match dist1 and dist2 and calculating per class ratio diff after adjustment # noqa: SC100
      dist1_counts, dist2_counts = _balance_sizes_downsizing(
          dist1_counts, dist2_counts, round_to_int=False
      )
      multipliers = [
          x2 / x1 for x1, x2 in zip(dist1_counts + 1, dist2_counts + 1)
      ]
  
      # changing dist1 to the even distribution and dist2 to be multiplication of dist1 by the multipliers # noqa: SC100
      dist1_counts = np.array(
          [int(np.sum(dist2_counts + 1) / len(dist2_counts))]
          * len(dist2_counts)
      )
      dist2_counts = np.round(dist1_counts * multipliers)
      dist2_counts = np.round(
          dist2_counts * (sum(dist1_counts) / sum(dist2_counts))
      )  # size normalization
      return dist1_counts, dist2_counts
  
imports: from numbers import Number
  from typing import Dict, Optional, Tuple, Union
  import numpy
  import pandas
  from plotly.graph_objs import Figure
  from plotly.subplots import make_subplots
  from scipy.stats import chi2_contingency, wasserstein_distance
  from deepchecks.core import ConditionCategory, ConditionResult
  from deepchecks.core.errors import DeepchecksValueError, NotEnoughSamplesError
  from deepchecks.utils.dict_funcs import get_dict_entry_by_value
  from deepchecks.utils.distribution.plot import CategoriesSortingKind, drift_score_bar_traces, feature_distribution_traces
  from deepchecks.utils.distribution.preprocessing import preprocess_2_cat_cols_to_same_bins
  from deepchecks.utils.plot import DEFAULT_DATASET_NAMES
  from deepchecks.utils.strings import format_number
label_desc: Use image node_id 3 for calling a global function with example usage: rebalance_distributions(dist1_counts, dist2_counts) and returns: dist1_counts,  dist2_counts