35. Information gain

medium

General

senior

Compute information gain for a candidate split in a decision tree, which tells you how much uncertainty about the labels is reduced after splitting. You’ll implement entropy and use it to score a split given arrays of labels and feature values.

The Information Gain for a binary split is defined as:

IG(Y) = H(Y) - \left( \frac{n_{left}}{n} H(Y_{left}) + \frac{n_{right}}{n} H(Y_{right}) \right)

where $H(Y)$ is the entropy of the labels.

Requirements

Implement the function

Python

def information_gain(y, x, threshold):
    """
    Computes information gain for a binary split x <= threshold vs x > threshold.

    Args:
        y (np.ndarray): Class labels (e.g., 0/1 or 0..K-1), length n.
        x (np.ndarray): Feature values aligned with y, length n.
        threshold (float): Split threshold.

    Returns:
        float: Information gain IG = H(y) - [ (n_L/n) * H(y_L) + (n_R/n) * H(y_R) ].
    """
    # Your solution below

Rules:

Use entropy $H(S) = -\sum_{c} p(c)\log_2 p(c)$ (ignore terms where $p(c)=0$ ).
Split into left group $L=\{i: x_i \le t\}$ and right group $R=\{i: x_i > t\}$ .
Use only NumPy.
Return a single float.

Example

Python

y = np.array([0, 0, 1, 1])
x = np.array([2.0, 3.0, 8.0, 9.0])
threshold = 5.0

information_gain(y, x, threshold)

Output:

Python

1.0

Function contract

Inputs

x: np.ndarray
y: np.ndarray
threshold: float

Returns

value: float

Examples

Input

{
  "x": [
    2,
    3,
    8,
    9
  ],
  "y": [
    0,
    0,
    1,
    1
  ],
  "threshold": 5
}

Output

Argument	Type
x	np.ndarray
y	np.ndarray
threshold	float

Return Name	Type
value	float

Constraints

Use NumPy.
Return a float.

Hint 1

Use np.unique(y, return_counts=True) to get class counts for entropy calculation.

Hint 2

Entropy is -sum(p * log2(p)). Handle p=0 implicitly by only iterating over non-zero counts.

Hint 3

Use boolean masking mask = x <= threshold to split y into y[mask] and y[~mask].

Roles

ML Engineer

AI Engineer

Data Scientist

Research Scientist

Companies

General

Levels

senior

entry

35. Information gain

Requirements

Example

Examples

Constraints

Similar Questions