Zenn
🎯

【Java】しゃくとり法

2024/05/01に公開
Java
初心者
tech

しゃくとり法とは

しゃくとり法とは、最適化問題や探索問題におけるアルゴリズムの一種です。ある配列の中から最適な値を、例えば、条件を満たす最長の区間や条件を満たす区間の個数等を探索する際に使用します。条件を満たす満たす場合にはしゃくを広げ、条件を満たさない場合はしゃくを縮めることにより最適解を求めます。

例題

サイズnの配列aの部分配列のうち、要素の合計がk以下となる部分配列の個数を求めよ。

例題の値

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

例題の答え

[2], [2, 3], [2, 3, 1], [3], [3, 1], [1], [1, 4], [4], [4, 2], [2]

の10個が答えとなる部分配列です。

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

  1. 右端のインデックスを表す変数をr, 現在の総和を表す変数をsum, 答えである個数を表す変数をcountで定義して、0で初期化します。
  2. 左端のインデックスを表す変数lを0からn - 1まで繰り返すfor文を定義します。
  3. 条件式が、rがnより小さく、かつ、sumにa[r]を加えてもkを超えないというwhile文を定義します。
  4. 条件式を満たす場合、sumにa[r]を加えて右端rを1進めます。
  5. 条件式を満たさない場合、while文を抜けて、countにr - lを加えます。
  6. 条件式が、l == rというif文を定義します。
  7. 条件式を満たす場合、右端のrを1進めます。
  8. 条件式を満たさない場合、sumからa[l]を引きます。
  9. 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, 現在の総和の変数を sum, 答えである個数の変数を count で定義し、0 で初期化
        int r = 0, sum = 0, count = 0;
        // 左端のインデックスの変数 l を 0 から n - 1 まで繰り返す
        for (int l = 0; l < n; l++) {
            // r が n より小さく、sum に a[r] を加えても k を超えない間
            while (r < n && sum + a[r] <= k) {
                // sum に a[r] を加えて右端 r を 1 進めます
                sum += a[r];
                r++;
            }
            // r - l を count に加える
            count += r - l;
            // もし l が r と同じなら
            if (l == r) {
                // r を 1 進めます
                r++;
            // そうでなければ
            } else {
                // sum から a[l] を引く
                sum -= a[l];
            }
        }
        // count を出力
        System.out.println(count);
    }
}

入力値

5 6
2 3 1 4 2

出力結果

10

Discussion