C++で範囲内の非ゼロの桁がK個以下となる数値の個数を求める方法

本記事では、整数の範囲(変数 start から end まで)と整数 k が与えられたとき、「ゼロ以外の数字(非ゼロの桁)が k 個以下しか含まれない数値」がその範囲内にいくつあるかを求める方法を解説します。この問題は「桁DP」と呼ばれる手法とメモ化再帰を組み合わせることで、大きな範囲でも効率的に解くことができます。

具体例

入力 - int start = 50, end = 100, K = 2;

出力 - 範囲内で非ゼロの桁がK個以下である数値の個数:50

説明 - 範囲は50から100、k は2です。50〜99の2桁の数値は、どれも非ゼロの桁が最大でも2個であるため条件を満たします。唯一の3桁の数値である100も、非ゼロの桁は「1」の1個だけなので条件を満たします。なお、本プログラムでは f(end) − f(start) という計算を行うため start 自身はカウント対象外となり、答えは50になります。

入力 - int start = 50, end = 100, K = 1;

出力 - 範囲内で非ゼロの桁がK個以下である数値の個数:5

説明 - 範囲は同じく50から100ですが、今度は k が1です。2桁の数値のうち非ゼロの桁が1個以下なのは、下一桁が0の数値だけです。f(end) − f(start) の計算により、該当するのは 60、70、80、90、そして100の5つとなるため、答えは5になります。

プログラムで採用しているアプローチ

• start から end までの整数範囲を用意し、k を宣言して値を入力します。入力データは後続の処理を行う関数へ渡されます。
• vector 型の変数(ここでは vec とします)を作成します。
• while ループを使い、start の値(val)が0になるまで、val % 10 を vec に追加し、val を val / 10 で更新していきます。これにより数値を1桁ずつ分解できます。
• STL の reverse() 関数に vec.begin() と vec.end() を引数として渡し、桁の並びを上位桁から下位桁の順に整えます。
• memset を使い、メモ化用の配列 arr の全要素を -1 で初期化します。
• check_val(0, 0, 0, vec) を返します。この関数が、各桁が非ゼロかどうかを判定しながら条件を満たす数値を数えます。
• check_val 関数の内部処理は次のとおりです。
  • place が vector のサイズと等しい場合、temp が k 以下であれば 1 を、そうでなければ 0 を返します。
  • arr[place][temp][set_val] が -1 以外の場合は、すでに計算済みのためその値をそのまま返します(メモ化による高速化)。
  • 結果を格納する変数を宣言し、0 で初期化します。
  • 変数 val を宣言し、set_val が 1 の場合は 9(以降の桁を自由に選べる)、そうでなければ vec[place](現在の桁の上限値)を設定します。
  • 0 から val まで for ループを実行します。
  • ループ内では、まず temp_2 に temp をコピーし、i が 0 でなければ(=非ゼロの桁を選んだ場合)temp_2 をインクリメントします。同様に temp_3 に set_val をコピーし、i が vec[place] より小さい場合は temp_3 を 1 に更新します(以降の桁が自由に選べる状態を表します)。
  • count に再帰呼び出し check_val(place + 1, temp_2, temp_3, vec) の結果を加算していきます。
  • 最後に arr[place][temp][set_val] = count を返します。

サンプルコード(C++)

#include <bits/stdc++.h>
using namespace std;
int arr[20][20][2];
int K;
int check_val(int place, int temp, int set_val, vector < int > vec) {
   if (place == vec.size()) {
      if (temp <= K) {
         return 1;
      }
      return 0;
   }
   if (arr[place][temp][set_val] != -1) {
      return arr[place][temp][set_val];
   }
   int count = 0;
   int val = (set_val ? 9 : vec[place]);
   for (int i = 0; i <= val; i++) {
      int temp_2 = temp;
      if (i != 0) {
         temp_2++;
      }
      int temp_3 = set_val;
      if (i < vec[place]) {
         temp_3 = 1;
      }
      count += check_val(place + 1, temp_2, temp_3, vec);
   }
   return arr[place][temp][set_val] = count;
}

int Not_more_k(int val) {
   vector < int > vec;
   while (val) {
      vec.push_back(val % 10);
      val = val / 10;
   }
   reverse(vec.begin(), vec.end());
   memset(arr, -1, sizeof(arr));
   return check_val(0, 0, 0, vec);
}
int main() {
   int start = 50, end = 100;
   K = 2;
   int count = Not_more_k(end) - Not_more_k(start);
   cout << "Count of Numbers in Range where the number does not contain more than K non zero digits are: " << count;
   return 0;
}

上記のコードを実行すると、次の出力が得られます。

出力

Count of Numbers in Range where the number does not contain more than K non zero digits are: 50

計算量について

このアルゴリズムの時間計算量は、桁数を d、非ゼロ桁の上限を K とすると O(d × K × 2 × 10) 程度に収まります。メモ化によって同じ状態の再計算が避けられるため、start と end の差が非常に大きい場合でも高速に動作するのが大きな特徴です。