问题描述

给定一个大小为N的整数数组A，你需要编写一个程序来计算数组A中每个元素与其右边“下一个更大元素”之间的最小距离。

函数签名

def min_distance_to_next_greater_element(arr: List[int]) -> List[int]:
    pass

参数说明

• arr: 长度为 $N$ 的整数数组 $A$
• 返回值: 长度为 $N$ 的整数数组，其中第 $i$ 个元素表示 $A[i]$ 与其右边 “下一个更大元素” 之间的最小距离。如果不存在 “下一个更大元素”，则返回 $-1$。

例子

Input: arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
Output: [1, 2, 2, 1, 3, 4, 1, 0, -1]
Explanation: 
arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
           ^--------^
           |     |
           6     8
           
arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
              ^-----^
              |  |
              3  8
              
arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
                 ^--^
                 |  |
                 9  10
                 
arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
                    ^--^
                    |  |
                    8  10
              
arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
                          ^--^
                          |  |
                          10 15

arr = [6, 3, 9, 8, 10, 2, 1, 15, 7]
                              ^---^
                              |   |
                              1   7

所以输出为 [1, 2, 2, 1, 3, 4, 1, 0, -1].

解法

本题的主要思路是使用栈维护每个元素的下一个更大元素。

具体地，我们从右往左扫描数组 $A$，对于每个元素 $A[i]$，我们将它和在栈中目前处于栈顶的元素 $A[stack[-1]]$ 进行比较。如果 $A[i] > A[stack[-1]]$，则 $A[stack[-1]]$ 的下一个更大元素就是 $A[i]$，并且此时 $A[i]$ 的最小距离为 $1$； 如果 $A[i] \leq A[stack[-1]]$，我们就一直弹出栈顶，直到栈为空或者 $A[i] > A[stack[-1]]$ 为止。

最后，如果弹出 $k$ 个元素后栈为空，则 $A[k]$ 的最小距离为 $-1$。

具体实现细节可以见下面的代码。

代码实现

from typing import List


def min_distance_to_next_greater_element(arr: List[int]) -> List[int]:
    n = len(arr)
    ans = [-1] * n
    stack = []

    for i in range(n - 1, -1, -1):
        while stack and arr[i] >= arr[stack[-1]]:
            stack.pop()

        if stack:
            ans[i] = stack[-1] - i
        
        stack.append(i)
    
    for i in range(n):
        if ans[i] != -1 and ans[i] != 1:
            for j in range(i + 1, i + ans[i]):
                if ans[j] == -1:
                    ans[j] = j - i
    
    return ans

时间复杂度

该算法的时间复杂度为 $O(n)$。