JZX轻语：简
LeetCode 2940 - 找到 Alice 和 Bob 可以相遇的建筑
发表于2024年08月10日
一开始自然地想到了单调栈：先假定bi > ai(如果bi == ai，已经相遇，则直接返回ai即可)，如果bi比ai高，则Alice直接走到bi即可；否则，Bob需要找到第一个比ai高的建筑。此类找到右侧第一个比当前位置高的建筑的问题，可以使用单调栈解决: 如果第一个比bi高的建筑(即为ci)也比ai高，则Alice可以直接走到这个建筑；否则，Bob需要继续找到右侧第一个比ai高的建筑: 类似于链表的形式，Bob继续跳到比ci高的第一个建筑，直到找到比ai高的建筑。这样的时间复杂度为O(n2)。且链表状的查找过程，无法使用二分查找优化。
因此，我们需要使用ST表/线段树/树状数组来优化此类的区间查询问题: 即在区间[bi, n-1]中找到第一个比ai高的建筑。这样的时间复杂度可降为O(nlogn)。
这里采用的是ST表的解法，ST表的单次查询时间复杂度为O(1)，但此类问题不能直接查询，需要逐步将查询区间缩小到一个点: 采用二分的形式，如果左半部分有满足条件的(即左半区间的最大值大于目标值，第一个比目标值大的元素必然在左区间)，缩小右边界；否则，满足条件的值只能在右半部分。这样的时间复杂度为O(logn)。
from typing import Callable, Generic, TypeVar
T = TypeVar('T')
class SparseTable(Generic[T]):
def __init__(self, list_: list[T], op_func: Callable[[T, T], T]):
"""
@param list_: 原始列表
@param op_func: 运算函数, 接受两个参数, 返回一个参数
"""
self._list = list_
self._op_func = op_func
elem_type = type(self._list[0]) if self._list else type(None)
l = len(self._list)
log2_l = math.floor(math.log2(l))
self._st = [[elem_type()] * (log2_l + 1) for _ in range(l)] # type: list[list[T]]
for i in range(l):
self._st[i][0] = self._list[i]
for j in range(1, log2_l + 1): # 1..log2_l
range_size = 1 << j
for i in range(l - range_size + 1):
# merge
self._st[i][j] = self._op_func(
self._st[i][j - 1],
self._st[i + (1 << (j - 1))][j - 1]
)
def query(self, l: int, r: int) -> T:
""" 查询区间 [l, r] 的结果 """
range_size = r - l + 1
sub_range = math.floor(math.log2(range_size))
return self._op_func(self._st[l][sub_range], self._st[r - (1 << sub_range) + 1][sub_range])
def __len__(self):
return len(self._list)
def __getitem__(self, item):
if not isinstance(item, tuple) or len(item) != 2:
raise TypeError('SparseTable.__getitem__ only accepts 2-tuple')
l, r = item
return self.query(l, r)
class Solution:
def leftmostBuildingQueries(self, heights: List[int], queries: List[List[int]]) -> List[int]:
n = len(heights)
st = SparseTable(list(range(n)), lambda a, b: a if heights[a] > heights[b] else b)
def find_next_higher(l: int, r: int, target: int) -> int:
""" 找到区间 [l, r] 中第一个高度大于等于 target 的位置 """
if heights[st.query(l, r)] <= target: # 如果区间最大值都小于 target, 直接返回 -1
return -1
while l < r:
mid = (l + r) >> 1
if heights[st.query(l, mid)] > target: # 如果左半部分有满足条件的, 缩小右边界
r = mid
else: # 否则，满足条件的值只能在右半部分
l = mid + 1
return l
ans = [-1] * len(queries)
for i, (u, v) in enumerate(queries):
if u == v:
ans[i] = v
continue
if u > v:
u, v = v, u
if heights[v] > heights[u]:
ans[i] = v
continue
ans[i] = find_next_higher(v, len(heights) - 1, heights[u])
return ans超时的版本，使用的是单调栈：
class Solution:
def leftmostBuildingQueries(self, heights: List[int], queries: List[List[int]]) -> List[int]:
next_higher = [-1] * len(heights) # next_higher[i] 表示 i 右侧第一个高度大于等于 heights[i] 的位置
stack = []
for i, height in enumerate(heights):
while stack and stack[-1][0] < height:
next_higher[stack[-1][1]] = i
stack.pop()
stack.append((height, i))
del stack
ans = [-1] * len(queries)
for i, (u, v) in enumerate(queries):
if u == v:
ans[i] = v
continue
if u > v:
u, v = v, u
while v != -1:
if heights[v] > heights[u]:
ans[i] = v
break
v = next_higher[v]
return ans