FastPrepDecision Tree Classifier from Scratch

Decision Tree Classifier from Scratch

Pinterest logoPinterest● HardFULLTIMEOA
Learn

Problem statement

Implement a decision tree classifier from scratch. You are given training rows xTrain, their numeric class labels yTrain, and unlabeled test rows xTest. Return the predicted class label for each test row.

An internal node stores a feature index and a threshold. A leaf stores a class label. To classify a row, follow the feature comparisons until reaching a leaf.

Entropy and information gain

For a nonempty set of rows, let p(c) be the fraction whose label is c. Its entropy is H = -sum(p(c) * log2(p(c))), summing only over labels present in the set. For a split into nonempty children L and R, the information gain is H(parent) - (|L| * H(L) + |R| * H(R)) / |parent|.

Rules for this exercise

The following rules make training and prediction deterministic:

  1. At each node, consider feature indices in ascending order. For each feature, sort its distinct values among the rows at that node. Consider the midpoint of every adjacent pair of distinct values, in ascending order, as a candidate threshold. Compute a midpoint as a + (b - a) / 2.
  2. Send rows with a feature value less than or equal to the threshold to the left child; send greater values to the right child. Ignore a threshold if floating-point rounding would leave either child empty. Positive and negative zero represent the same numeric value.
  3. Start the best gain at zero. Replace the selected split only when a candidate gain is greater than the current best gain by more than 10^-12. Thus a tied gain keeps the earlier feature and then the earlier threshold.
  4. If every row has the same label, create a leaf immediately. Also create a leaf when no candidate improves the initial zero gain by more than 10^-12. A leaf uses the most frequent label at that node; equal frequencies choose the smallest numeric label.
  5. Otherwise, recursively train both children using the selected split. There is no depth limit or additional pruning.
  6. Return labels in the original order of xTest. An empty test set returns an empty array. Labels are numeric categories and may be negative or fractional; do not round or convert them to integers.

Implement the entropy, split selection, tree construction, and prediction logic yourself. Standard containers, sorting, and logarithm functions may be used; prebuilt machine-learning implementations may not.

Function

decisionTreePredictions(xTrain: double[][], yTrain: double[], xTest: double[][]) → double[]

Examples

Example 1

xTrain = [[0.0],[1.0],[3.0],[4.0]]yTrain = [1.5,1.5,2.5,2.5]xTest = [[2.0],[2.1],[-5.0],[9.0]]return = [1.5,2.5,1.5,2.5]

The midpoint 2 separates the two labels into pure children. A test value equal to the threshold goes left. The returned labels follow the order of xTest.

Example 2

xTrain = [[1.0,1.0],[1.0,1.0]]yTrain = [-0.5,2.25]xTest = [[0.0,0.0],[2.0,2.0]]return = [-0.5,-0.5]

Both training rows have identical features, so no threshold exists. Their label counts tie, and the leaf predicts the smaller numeric label, -0.5, for both test rows.

Constraints

  • 1 <= xTrain.length == yTrain.length <= 80.
  • Each training row has 1 to 6 features. Every training and test row has the same feature count.
  • 0 <= xTest.length <= 80.
  • All inputs are finite floating-point numbers. Feature values are between -10^6 and 10^6, inclusive; labels are between -1000 and 1000, inclusive.
  • Training rows may repeat, including with different class labels.

More Pinterest problems

See Pinterest hiring insights
public double[] decisionTreePredictions(double[][] xTrain, double[] yTrain, double[][] xTest) {
    // Write your code here.
}
xTrain[[0.0],[1.0],[3.0],[4.0]]
yTrain[1.5,1.5,2.5,2.5]
xTest[[2.0],[2.1],[-5.0],[9.0]]
expected[1.5,2.5,1.5,2.5]
Checking account…