Easy Understood TreeMap Solution

  • 2

    The squares divide the number line into many segments with different heights. Therefore we can use a TreeMap to store the number line. The key is the starting point of each segment and the value is the height of the segment. For every new falling square (s, l), we update those segments between s and s + l.

    class Solution {
        public List<Integer> fallingSquares(int[][] positions) {
            List<Integer> list = new ArrayList<>();
            TreeMap<Integer, Integer> map = new TreeMap<>();
            // at first, there is only one segment starting from 0 with height 0
            map.put(0, 0);
            // The global max height is 0
            int max = 0;
            for(int[] position : positions) {
                // the new segment 
                int start = position[0], end = start + position[1];
                // find the height among this range
                Integer key = map.floorKey(start);
                int h = map.get(key);
                key = map.higherKey(key);
                while(key != null && key < end) {
                    h = Math.max(h, map.get(key));
                    key = map.higherKey(key);
                h += position[1];
                // update global max height
                max = Math.max(max, h);
                // update new segment and delete previous segments among the range
                int tail = map.floorEntry(end).getValue();
                map.put(start, h);
                map.put(end, tail);
                key = map.higherKey(start);
                while(key != null && key < end) {
                    key = map.higherKey(key);
            return list;

  • 0

    I think in line
    int tail = map.lowerEntry(end).getValue();
    the lowerEntry should actually be floorEntry, right?

  • 0

    @ywsstan Thank you for pointing out. You are right. It should be the floorEntry. It seems such test case is missing.

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.