C++で最後のK個の数値の積を効率的に求める方法

問題概要

本記事では、ProductOfNumbersというクラスを実装する方法を解説します。このクラスは、以下の2つのメソッドをサポートする必要があります。

• add(int num): 現在の数値リストの末尾に数値 num を追加します。
• getProduct(int k): 現在のリスト内の末尾 k 個の数値の積を返します。

なお、getProduct が呼び出される時点で、現在のリストには必ず k 個以上の数値が存在すると仮定して構いません。

動作例

例えば、以下の順序でメソッドを呼び出した場合を考えてみましょう。

add(3) → [3]
add(0) → [3, 0]
add(2) → [3, 0, 2]
add(5) → [3, 0, 2, 5]
add(4) → [3, 0, 2, 5, 4]
getProduct(2) → (5 × 4) = 20
getProduct(3) → (2 × 5 × 4) = 40
getProduct(4) → (0 × 2 × 5 × 4) = 0
add(8) → [3, 0, 2, 5, 4, 8]
getProduct(2) → (4 × 8) = 32

解法のアプローチ：累積積（プレフィックス積）の活用

この問題を効率的に解く鍵となるのが「累積積」の考え方です。毎回末尾 k 個を掛け直す代わりに、追加するたびにそれまでの総積を記録しておくことで、積の取得を高速に行えます。具体的な手順は以下の通りです。

• 初期化時: 配列を作成し、初期値として 1 を格納します。
• add(num) の処理: もし num が 0 ならば、配列をクリアして 1 を挿入します。それ以外の場合は、「末尾の要素 × num」を配列に追加します。
• getProduct(k) の処理: 配列のサイズを n としたとき、k > n − 1 であれば 0 を返します。そうでなければ dp[n - 1] / dp[n - k - 1] を返します。

ここで重要なのは、0 が出現した時点でリストをリセットする点です。一度 0 が加わると、それ以降の積は必ず 0 になるため、配列をクリアして計算をやり直すことで、除算を使った積の算出が常に正しく機能するようになります。

C++での実装例

それでは、実際のコードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class ProductOfNumbers {
public:
    vector<int> dq;
    ProductOfNumbers() {
        dq.push_back(1);
    }
    void add(int num) {
        if(num == 0){
            dq.clear();
            dq.push_back(1);
        }
        else{
            dq.push_back(dq.back() * num);
        }
    }
    int getProduct(int k) {
        int n = (int)dq.size();
        return k > n - 1 ? 0 : dq[n - 1] / dq[n - k - 1];
    }
};
main(){
    ProductOfNumbers ob;
    (ob.add(3));
    (ob.add(0));
    (ob.add(2));
    (ob.add(5));
    (ob.add(4));
    cout << (ob.getProduct(2)) << endl;
    cout << (ob.getProduct(3)) << endl;
    cout << (ob.getProduct(4)) << endl;
    (ob.add(8));
    cout << (ob.getProduct(2)) << endl;
}

入力

add(3)
add(0)
add(2)
add(5)
add(4)
getProduct(2)
getProduct(3)
getProduct(4)
add(8)
getProduct(2)

出力

20
40
0
32

計算量について

この実装の優れた点は、どちらの操作も定数時間で完了することです。

• add 操作: O(1) — 配列の末尾への要素追加のみで済みます。
• getProduct 操作: O(1) — あらかじめ記録しておいた累積積の除算 1 回だけで結果が得られます。

素朴な実装では getProduct(k) のたびに最大 k 回の乗算が必要となり O(k) の計算量がかかりますが、累積積を用いることで大幅な高速化が実現できます。数値の追加と積の参照が頻繁に行われる場面で特に有効な手法です。