C++でインクリメント操作を効率的に扱うスタックの設計方法

本記事では、以下の操作をサポートするカスタムスタックをC++で設計する方法を解説します。ポイントは、インクリメント操作を毎回全要素に適用するのではなく、遅延評価（lazy propagation）の考え方を使って効率化することです。

サポートする操作

• CustomStack(int maxSize)：最大要素数 maxSize を持つスタックを初期化します。スタックが maxSize に達した場合、それ以上の追加は行われません。
• void push(int x)：スタックのサイズが maxSize に達していない場合、要素 x をスタックの先頭に挿入します。
• int pop()：スタックの先頭要素を削除して返します。スタックが空の場合は -1 を返します。
• void increment(int k, int val)：スタックの下から k 個の要素に val を加算します。スタック内の要素が k 個未満の場合は、すべての要素に加算します。

解法のアプローチ

increment 操作のたびに下から k 個の要素を直接更新すると、計算量が O(k) かかり非効率です。そこで、加算値を記録する補助配列を用意し、pop のタイミングでまとめて反映させる手法を採用します。

手順

1. スタック本体の配列 st、加算値を記録する配列 inc、容量を表す整数型変数 cap を定義します。
2. コンストラクタで cap := N とし、inc をサイズ N + 10 の配列として初期化します。
3. push(x)：スタックのサイズが cap でない場合のみ、x を st に挿入します。
4. pop()：以下のように動作します。
   • st が空なら -1 を返します。
   • そうでなければ：
     • スタックの最上位要素に inc[最上位のインデックス] を加算します。
     • スタックに要素が残っている場合、inc[サイズ - 2] に inc[サイズ - 1] を加算します（下位の要素へ加算値を引き継ぐ）。
     • inc[サイズ - 1] を 0 にリセットします。
     • st の末尾要素を取り出し、その値を返します。
5. increment(k, val)：以下のように動作します。
   • k を 1 減らします（0始まりのインデックスに変換）。
   • k := min(k, スタックのサイズ - 1) とします。
   • k < 0 の場合は何もしません。
   • inc[k] に val を加算します。

この方法により、push・pop・increment はいずれも O(1) の計算量で実行できます。

C++での実装例

以下の実装を見ると、より理解が深まります。

#include <bits/stdc++.h>
using namespace std;
class CustomStack {
public:
   vector <int> st;
   vector <int> inc;
   int cap;
   CustomStack(int N) {
      cap = N;
      inc = vector <int>(N + 10);
   }
   void push(int x) {
      if(st.size() == cap) return;
      st.push_back(x);
   }
   int pop() {
      if(st.empty()) return -1;
      else{
         st.back() += inc[st.size() - 1];
         if(st.size() - 1 > 0 ){
            inc[st.size() - 2] += inc[st.size() - 1];
         }
         inc[st.size() - 1] = 0;
         int x = st.back();
         st.pop_back();
         return x;
      }
   }
   void increment(int k, int val) {
      k--;
      k = min(k, (int)st.size() - 1);
      if(k < 0) return;
      inc[k] += val;
   }
};
main(){
   CustomStack ob(3);
   ob.push(1);
   ob.push(2);
   cout << ob.pop() << endl;
   ob.push(2);
   ob.push(3);
   ob.push(4);
   ob.increment(5, 100);
   ob.increment(2, 100);
   cout << ob.pop() << endl;
   cout << ob.pop() << endl;
   cout << ob.pop() << endl;
   cout << ob.pop() << endl;
}

入力

プログラム内の main() 関数を参照してください

出力

2
103
202
201
-1

動作の解説

まず maxSize = 3 のスタックを作成し、1 と 2 を push した後、pop を呼び出すと 2 が返されます。続いて 2、3、4 を push しようとしますが、最大サイズが 3 のため 4 は無視されます。increment(5, 100) では要素数が 5 未満のため全要素（2 と 3）に 100 が加算対象となり、increment(2, 100) では下から 2 要素に 100 が加算されます。以降の pop の結果として 103、202、201 が順に返され、最後に空のスタックに対する pop で -1 が出力されます。