【Java】しゃくとり法の応用

例題

昇順にソートされた長さnの配列aがあります。
配列の要素a[i]に整数kを与えた場合のa[i] + k以下となる配列aの要素の個数を求めよ。

例題の値

n = 5
k = 2
a = [1, 2, 3, 4, 5]

例題の答え

i = 0 の時
a[0] + k = 1 + 2 = 3
3以下となる要素の個数は[1,2,3]の3個

i = 1 の時
a[1] + k = 2 + 2 = 4
4以下となる要素の個数は[1,2,3,4]の4個

i = 2 の時
a[2] + k = 3 + 2 = 5
5以下となる要素の個数は[1,2,3,4,5]の5個

i = 3 の時
a[3] + k = 4 + 2 = 6
6以下となる要素の個数は[1,2,3,4,5]の5個

i = 4 の時
a[4] + k = 5 + 2 = 7
7以下となる要素の個数は[1,2,3,4,5]の5個

例題におけるしゃくとり法のアルゴリズム

1. 右端のインデックスを表す変数をrを0で初期化します。
2. 左端のインデックスを表す変数lを0からn - 1まで繰り返すfor文を定義します。
3. 条件式が、rがnより小さく、かつ、a[l]にkを加えてもa[r]を超えないというwhile文を定義します。
4. 条件式を満たす場合、右端rを1進めます。
5. 条件式を満たさない場合、rを出力します。

   l = 0 の時
   a[l] + k = 1 + 2 = 3
a[r] = a[0] = 1
1 <= 3 は、条件式を満たすので、rを1進めます。
   a[r] = a[1] = 2
2 <= 3 は、条件式を満たすので、rを1進めます。
   a[r] = a[2] = 3
3 <= 3 は、条件式を満たすので、rを1進めます。
   a[r] = a[3] = 4
4 <= 3 は、条件式を満たさないので、r = 3を出力します。
   ※例題の答えと一致します。

6. 条件式が、l == rというif文を定義します。
7. 条件式を満たす場合、右端のrを1進めます。
8. 条件式を満たさない場合、lを1進めた状態で、2.以降のfor文に戻ります。

コード例

import java.util.*;

public class Main {

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n, k;
        n = sc.nextInt();
        k = sc.nextInt();

        int[] a = new int[n];
        for (int i = 0; i < n; i++) {
            a[i] = sc.nextInt();
        }

        // 右端のインデックスの変数を r を 0 で初期化
        int r = 0;
        // 左端のインデックスの変数 l を 0 から n - 1 まで繰り返す
        for (int l = 0; l < n; l++) {
            // r が n より小さく、a[l] に k を加えても a[r] を超えない間
            while (r < n && a[r] <= a[l] + k) {
                // 右端 r を 1 進めます
                r++;
            }

            // r を出力
            System.out.println(r);

            // もし l が r と同じなら
            if (l == r) {
                // r を 1 進めます
                r++;
            }
        }
    }
}

入力値

5 2
1 2 3 4 5

出力結果

3
4
5
5
5