P2650. 找朋友

解题思路

两种写法的比较

注意，本题和 每日温度 非常类似。区别在于，本题需要找到的是 右边下一个更大元素的索引，而非与当前元素的间隔，显然变得更加简单了。

我们讲过，类似这种要求寻找 左边/右边最近的更大/更小元素 的题目，均可以使用 单调栈 来完成。

对于单调栈的题目，既可以 正序遍历 也可以 逆序遍历 数组来完成，重点在于理解单调栈的原理，同学们只需要选择适合自己理解的方法来完成即可。以下表格总结了两种不同遍历顺序的异同点。

---

正序遍历图解

正序遍历原数组，初始化单调栈为空，答案数组均为 0。

1. 遍历到 nums[0]，单调栈中没有元素，索引 0 入栈。 2. 继续遍历到 nums[1]，栈顶索引 0 对应的 123 小于当前元素 124。 栈顶元素 0 出栈，同时修改 ans[0] 为当前索引 1。 出栈和修改结束后，当前索引 1 入栈。 对于正序遍历写法而言：修改答案是在出栈的时候进行的。 3. 遍历到 nums[2]，过程也是类似的。 4. 继续遍历 nums[3] 和 nums[4]，由于 121 和 119 相较于 125 都是递减的，因此直接入栈，无需修改 ans。 5. 继续遍历 nums[5]，当前元素大于栈顶索引 4 和 3 对应的元素的 119 和 121。 将它们弹出，同时依次修改 ans[4] 和 ans[3] 为当前索引 5。 出栈和修改结束后，当前索引 5 入栈。 6. 继续遍历 nums[6]，当前元素大于栈顶索引 5 和 2 对应的元素的 122 和 125。 将它们弹出，同时依次修改 ans[5] 和 ans[2] 为当前索引 6。 出栈和修改结束后，当前索引 6 入栈。 7. 最后遍历到 nums[7]，小于栈顶索引对应的元素 126，无需弹出也无需修改答案数组。直接入栈。

最终答案数组的结果即为答案。

因此 正序遍历原数组的单调栈算法 的整体框架如下：

---

逆序遍历图解

逆序遍历原数组，初始化单调栈为空，答案数组均为 0。

1. 遍历到 nums[7]，单调栈中没有元素，索引 7 入栈。 2. 继续遍历到 nums[6]，栈顶索引 7 对应的 123 小于等于当前元素 126。 即使遍历到 6 更前面的位置（比如索引 5），126 必然会挡住后面的这个 123，故 123 此时已经无用了。 逆序遍历写法的出栈操作其实隐含了贪心思想 栈中的索引 7 出栈。出栈后栈中没有元素，说明不存在位于右边的比 126 更大的值，无需修改答案数组，直接将当前索引 6 入栈。 3. 继续遍历到 nums[5]。由于 122 小于栈顶元素对应的 126，可知此时栈顶索引 6 就是当前元素 122 右边第一个最大的元素的索引，修改 ans[5]。 4. 同理，继续遍历到 nums[4]。由于 119 小于栈顶元素对应的 122，可知此时栈顶索引 5 就是当前元素 119 右边第一个最大的元素的索引，修改 ans[4]。而后将索引 4 入栈。 5. 继续遍历到 nums[3]。通过 while 循环，将栈顶所有对应元素小于等于当前元素 121 的索引弹出（即 119）。 全部弹出后，此时栈顶仍然存在索引 5，即为当前元素 121 右边的第一个更大值对应的索引。 修改 ans[3]。而后将索引 3 入栈。 6. 继续遍历到 nums[2]。通过 while 循环，将栈顶所有对应元素小于等于当前元素 125 的索引弹出（即 121、122）。 全部弹出后，此时栈顶仍然存在索引 6，即为当前元素 125 右边的第一个更大值对应的索引。 修改 ans[3]。而后将索引 2 入栈。 7. 继续依次遍历 nums[1] 和 nums[0]，其实过程和遍历到 nums[5] 和 nums[4] 是类似的，修改 ans[1] 和 ans[0]，最终得到答案数组。

因此 逆序遍历原数组的单调栈算法 的整体框架如下：

---

思考

作为一道经典的单调栈题目，我们现在已经掌握了其 正序 和 逆序 两种写法。 这道题是非常值得反复琢磨的题目。

请课后或者完成本题后再次思考以下问题：

• 为什么正序写法 while 的条件是 h > height[stack[-1]]，而逆序写法的是 h >= height[stack[-1]]？
• 如果改成找 右边第一个更小元素，代码应该如何修改？
• 如果改成找 左边第一个更大元素，代码应该如何修改？
• 如果改成找 左边第一个更小元素，代码应该如何修改？

复杂度分析 设同学人数为 n（题面约束 n 最大为 40000）。

• 时间复杂度：O(n)。虽然 for 循环里嵌套了一个 while，但要注意：每个下标 i 只会被 push 进栈一次；入栈之后，它至多在后面遇到第一个更高身高时被弹出一次。因此整个程序运行期间，入栈总次数和出栈总次数都不超过 n，while 的总代价是摊还线性的，并不是表面上的 O(n²)。读入身高数组和最后输出 ans 也都是 O(n)。
• 空间复杂度：O(n)。最坏情况是身高从左到右整体递减，此时没有任何下标被弹出，单调栈会存下全部 n 个下标；答案数组 ans 同样占 O(n)。

对比暴力做法（每个人向右逐个找第一个更高的人，最坏 O(n²)），单调栈的关键收益就在于「每个下标只进出栈一次」，在 n = 40000 的规模下轻松通过。