単調スタック:「次に大きい要素」を O(n²) から O(n) へ

各添字は高々 1 回 push され 1 回 pop されるので、二重ループに見える形が線形になります。左右で最初に大きい/小さいものを探す形なら、同じ雛形が使えます。

各要素について、右側で最初に自分より大きい値を探す。素直に書けば二重ループです。単調スタックなら 1 回の走査で済みます。代償は、スタックにまだ答えが出ていない添字が積まれると認めることです。

雛形

左から右へ走査し、スタックは減少列に保ちます。新しい要素が来たら、それより小さいものをすべて pop し、pop された要素は答えを得ます。

function nextGreater(nums) {
  const result = new Array(nums.length).fill(-1);
  const stack = [];
  for (let i = 0; i < nums.length; i += 1) {
    while (stack.length && nums[stack[stack.length - 1]] < nums[i]) {
      result[stack.pop()] = nums[i];
    }
    stack.push(i);
  }
  return result;
}

要点は、while の中で pop された要素がその場で答えを得て、以後二度と話題にならないことです。したがって内側の while は入れ子に見えても、合計の反復回数は O(n) です。

なぜ O(n) か

償却の議論は簡単です。各添字は for で 1 回 push され、pop は高々 1 回。合計の操作は 2n 以下で、データの並びに依存しません。これが、各要素から前方を走査する方法に勝つ唯一の理由です。後者は最悪の場合(減少列)毎回最後まで見に行きます。

同じ形の問題が 4 つ

問題 スタックが保つ順序 pop する条件
次の大きい要素 減少 新しい値が大きい
ヒストグラムの最大長方形 増加 新しい棒が低い
雨水を溜める 減少(高さ) 新しい棒が高い
毎日の気温 減少(添字) 新しい気温が高い

「左右で最初に大きい/小さいもの」と気づけば雛形が当てはまります。雨水を溜める問題は、pop しながら水たまりの面積を計算する変形です。

間違えやすい 2 点

  • 値ではなく添字を積む:距離や位置を問う問題(何日後に暖かくなるか)は値だけでは答えられません。添字なら常に安全です。
  • 等しいとき pop するか:< か <= かで、重複した要素が「厳密に大きい」隣か「以上」の隣を得るかが決まります。結果のどの位置が採用されるかが変わるので、書く前に意味を決めます。

左または右で最初に大きい/小さいものを問う問題なら、二重ループより先にスタックを考える。

← 記事一覧に戻る

コメント

…