用 Java 的 TreeMap(TreeSet) 配合哨兵实现 C++(lowerBound) 和 python(bisect_left) 中的二分查找。

参考:https://leetcode.cn/circle/discuss/xL7JQq/

可以通过 TreeSet 中的 lower、floor、ceiling、higher 实现C++ stl 中的 lower_bound 和 upper_bound。

TreeSet

// 返回此set中小于给定元素的最大元素,如果没有这样的元素,则返回NULL。
E lower(E e);

// 返回此集合中小于或等于给定元素的最大元素,如果没有这样的元素,则返回NULL。
E floor(E e);

// 返回此set中大于或等于给定元素的最小元素,如果没有这样的元素,则返回NULL。
E ceiling(E e);

// 返回此set中大于给定元素的最小元素,如果没有这样的元素,则返回NULL。
E higher(E e);

TreeMap中对应有

Map.Entry<K,V>	lowerEntry(K key)
Map.Entry<K,V>	floorEntry(K key)
Map.Entry<K,V>	ceilingEntry(K key)
Map.Entry<K,V>	higherEntry(K key)

Code

TreeMap 的 ceilingEntry 实现

class Solution {
    public List<Long> minOperations(int[] nums, int[] queries) {
        Arrays.sort(nums);
        int n = nums.length;
        var sum = new long[n + 1]; // 前缀和
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + nums[i];

        var ans = new ArrayList<Long>(queries.length);

        TreeMap<Integer, Integer> map = new TreeMap<>();
        for(int i = 0; i < n; i++){
            if(!map.containsKey(nums[i])){  // 注意重复元素只记录第一次出现的位置
                map.put(nums[i], i);
            }
        }
        map.put(Integer.MAX_VALUE, n);  // 添加哨兵
        for (int q : queries) {
            int j = map.ceilingEntry(q).getValue();  // 相当于 lowerBound
            long left = (long) q * j - sum[j]; // 蓝色面积
            long right = sum[n] - sum[j] - (long) q * (n - j); // 绿色面积
            ans.add(left + right);
        }
        return ans;
    }
}

自己实现 lowerBound

class Solution {
    public List<Long> minOperations(int[] nums, int[] queries) {
        Arrays.sort(nums);
        int n = nums.length;
        var sum = new long[n + 1]; // 前缀和
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + nums[i];

        var ans = new ArrayList<Long>(queries.length);
        for (int q : queries) {
            int j = lowerBound(nums, q);
            long left = (long) q * j - sum[j]; // 蓝色面积
            long right = sum[n] - sum[j] - (long) q * (n - j); // 绿色面积
            ans.add(left + right);
        }
        return ans;
    }

    // 见 https://www.bilibili.com/video/BV1AP41137w7/
    private int lowerBound(int[] nums, int target) {
        int left = -1, right = nums.length; // 开区间 (left, right)
        while (left + 1 < right) { // 区间不为空
            // 循环不变量:
            // nums[left] < target
            // nums[right] >= target
            int mid = left + (right - left) / 2;
            if (nums[mid] < target)
                left = mid; // 范围缩小到 (mid, right)
            else
                right = mid; // 范围缩小到 (left, mid)
        }
        return right;
    }
}