The best first question
A decision tree begins with one question “x ≤ t?”. There are n examples: each has a numeric feature xi and an answer yi, 0 or 1. The question splits the examples into two non-empty groups, and a good question is one after which uncertainty is smallest. Find the threshold t with the largest information gain
$$\mathrm{IG} = H(\text{all}) - \frac{n_{\le}}{n}\,H(\text{group } x \le t) - \frac{n_{>}}{n}\,H(\text{group } x > t),$$where H is the entropy of a group in bits: if the share of ones in a group is p, then H = −p log₂ p − (1 − p) log₂ (1 − p) (a term with a zero share counts as zero).
Two tests are earthquakes of magnitude 5 and above in 2015–2021 from the USGS catalogue: the feature is magnitude (in hundredths) or the distance to the nearest town of 10,000 people (km), the answer is whether anybody reported feeling the shock.
Input
The first line holds n (2 ≤ n ≤ 200,000). Each of the following n lines holds integers xi (|xi| ≤ 10⁹) and yi ∈ {0, 1}.
Output
If all xi are equal, print -: no question splits the examples. Otherwise print two numbers: the threshold t and the information gain. Any threshold will do at which both groups are non-empty and the gain differs from the largest possible by at most 10⁻⁹; the gain is checked with a tolerance of 10⁻⁶.
Examples
Example 1
8 1 0 2 0 3 0 4 1 5 0 6 1 7 1 8 1
3.5 0.548794941
Example 2
3 42 1 42 0 42 1
-
Limits
- Time · Python
- 6000 ms
- Time · Go
- 1500 ms
- Memory
- 256 MB
- Judging
- 8 tests
Walkthrough
Open it step by step: each step gives away a little more. You can jump straight to the solution, but a problem solved with one hint stays with you longer.
1Hint 1
There are few thresholds worth checking: the question “x ≤ t?” splits the examples the same way for all t between two adjacent distinct values of x. How many different splits can n numbers have?
2Hint 2
For each split you need only the sizes of the groups and the number of ones in each. If the examples are sorted by x, how do you get these numbers for the next split from the previous one in O(1)?
3The idea
Sort the examples by x and go from left to right, moving one example at a time from the right group to the left and updating the number of ones on the left. Each time adjacent values differ, we have one of the possible splits: compute its information gain and remember the best. The threshold is the midpoint between the adjacent values; any number between them will do as well.
4Walkthrough
The entropy of all the examples does not depend on the question, so maximising the gain is the same as minimising the average entropy of the groups weighted by their sizes. The weights are the main thing here: a group of one example is always “pure”, its entropy zero, but it describes one example in a thousand.
Equal x cannot be separated: the question “x ≤ t?” sends them all to one side. So a split between equal values is skipped, otherwise the program computes the gain of a split that cannot occur. If all x are equal, there are no splits at all.
On the earthquakes of 2015–2021 the best question about distance is “is the town no further than 106.5 km?” with a gain of 0.172 bits, and the best question about magnitude is “is the magnitude at most 5.85?” with a gain of only 0.042. That is why the tree in the article starts with distance.
4+Solution
5Complexity
O(n log n) for the sort and O(n) for the pass. Trying thresholds while recounting the groups from scratch is O(n²) and does not finish in time on two hundred thousand examples.
6Common mistakes
- An unweighted average of entropies. Cutting off an extreme example gives a pure group with zero entropy, and such a “question” wins. The seventh test is built on this: a lone example at the far edge.
- A split between equal values. The gain computed belongs to an impossible split, and the threshold printed gives a different one.
- log(0). An empty share in the entropy counts as zero: −0·log₂0 = 0 by continuity, not a computation error.