LeetCode 2940 找到Alice和Bob可以相遇的建筑

方法1：线段树

class Solution {
    static int[] tree; // 存储区间的最大值
    
    static void build(int o, int left, int right, int[] datas) {
        if(left == right) {
            tree[o] = datas[left - 1];
            return ;
        }
        int mid = (left + right) >> 1;
        build(o << 1, left, mid, datas);
        build(o << 1 | 1, mid + 1, right, datas);
        tree[o] = Math.max(tree[o << 1], tree[o << 1 | 1]);
    } 

    static int query(int start, int val, int o, int left, int right) {
        if(val >= tree[o]) // 当前区间所有值都不大于 val
            return 0; // 直接返回 0 出去后会被 -1
        if(left == right) // 前面条件没有退出 即当前值是第一个满足条件的
            return left;
        int mid = (left + right) >> 1;
        if(start <= mid) { // 当前位置在要求区间内
            int res = query(start, val, o << 1, left, mid);
            if(res > 0) return res;
        }
        // 当前位置不在要求区间内
        return query(start, val, o << 1 | 1, mid + 1, right);
    }

    public int[] leftmostBuildingQueries(int[] heights, int[][] queries) {
        int n = heights.length;
        tree = new int[n << 2];
        build(1, 1, n, heights);  // 建树
        int m = queries.length;
        int[] ans = new int[m];
        for(int i = 0; i < m; i++) {
            int a = queries[i][0];
            int b = queries[i][1];
            if(b < a) { // 保证 b 在 a 的右边
                int temp = a; a = b; b = temp;
            }
            // 在同一位置或者右边的位置比左边高
            if(a == b || heights[a] < heights[b]) {
                ans[i] = b; continue;
            }
            // 查询区间 [b + 1, n] 中高度大于 heights[a] 的下标
            ans[i] = query(b + 1, heights[a], 1, 1, n) - 1;
        }
        return ans;
    }
}

class Solution:
    def leftmostBuildingQueries(self, heights: List[int], queries: List[List[int]]) -> List[int]:
        n = len(heights); tree = [0] * (n << 2)
        
        def build(o: int, left: int, right: int) -> None:
            if left == right:
                tree[o] = heights[left - 1]
                return
            mid = (left + right) >> 1
            build(o << 1, left, mid)
            build(o << 1 | 1, mid + 1, right)
            tree[o] = max(tree[o << 1], tree[o << 1 | 1])
        
        def query(start: int, val: int, o: int, left: int, right: int) -> int:
            if val >= tree[o]:
                return 0
            if left == right:
                return left
            mid = (left + right) >> 1
            if start <= mid:
                res = query(start, val, o << 1, left, mid)
                if res > 0: return res
            return query(start, val, o << 1 | 1, mid + 1, right)
        
        build(1, 1, n); ans = list()
        for i, (a, b) in enumerate(queries):
            if a > b:
                a, b = b, a
            if a == b or heights[b] > heights[a]:
                ans.append(b); continue
            ans.append(query(b + 1, heights[a], 1, 1, n) - 1)
        return ans