Help us understand the problem. What is going on with this article?

【解答編】数列の左からの和がK以上となる最小の添字は?

More than 3 years have passed since last update.

問題編はこちら
この記事では↑の問題の解法を2通り、コードとともに紹介します。

解法1: 累積和×二分探索

もし N,Q が小さければ、クエリごとに毎回数列の上でループを回し、和が初めてK以上となったところで添え字を出力すればよいでしょう。これだと計算量はO(NQ)になります。しかし今回はN,Qの上限が10^5と大きいので、このアルゴリズムでは間に合いません。そこで、累積和二分探索を用います。

まず、数列aの1番目からi番目までの和(累積和)をすべて計算しておきます。
するとa_i>=0という制約から、aの累積和は単調増加になります。
単調性があるので二分探索ができます。
そこでクエリごとに、ループを回すのではなく二分探索で答えを探すと、計算量はO(N + QlogN)となります。(前処理にO(N), クエリ1つにO(logN))

ただしこの解法は、a_iが非負とは限らない場合には適用できないことに注意してください。

nibutan.cpp
#include <bits/stdc++.h>
using namespace std;
#define rep(i,n) for (int i=0;i<(n);i++)
#define rep2(i,a,b) for (int i=(a);i<(b);i++)
#define all(a) (a).begin(),(a).end()
typedef long long ll;
const ll INF = 1e18;

signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(0);

    int N, Q;
    while (cin >> N >> Q, N) {
        vector<ll> a(N + 2);
        rep2(i, 1, N + 1) {
            cin >> a[i];
            a[i] += a[i - 1];
        }
        a[0] = -INF;
        a.back() = INF;
        rep(i, Q) {
            ll K;
            cin >> K;
            int ans = lower_bound(all(a), K) - a.begin();
            cout << (ans == N + 1 ? -1 : ans)<< endl;
        }
    }
}

解法2: クエリ先読み

実はクエリ先読みをすると、a_iが負になりうる場合にも同等の計算量で対応できます
手順は以下の通りです。

  1. 解法1と同様に累積和をもっておく
  2. クエリをひとまずすべて読み込む。$q[id] = (K_{id}, id)$ のペアを $K_{id}$ の昇順にソートする
  3. 数列の累積和について、左からループを回す。累積和の添字を $i$, ソートされたクエリの添字を $j$ とし、答えは $ans[q[j].id] = (クエリidに対する答え)$ に記録するものとする。各 $i$ に対し、$j < Q$ かつ $a_i >= q[j].K$ 以上である限り $ans[q[j].id] = i, j++$ を繰り返す
  4. $ans$ を出力する

クエリをソートすることで、Kの小さい方から順に印をつけていくことができるというわけです。
計算量はO(N + QlogQ)となります。(前処理にO(N), クエリのソートにO(QlogQ), ↑のステップ3にO(N + Q))
この解法は累積和の単調性を仮定していないので、a_iが非負という制約がなくても有効です。

sakiyomi.cpp
signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(0);

    int N, Q;
    while (cin >> N >> Q, N) {
        vector<ll> a(N);
        rep(i, N) {
            cin >> a[i];
            if (i != 0) a[i] += a[i - 1];
        }

        vector<pair<ll, int>> queries;
        rep(id, Q) {
            ll K;
            cin >> K;
            queries.emplace_back(K, id);
        }
        sort(all(queries));

        vector<int> ans(Q, -2);
        int j = 0;
        rep(i, N) {
            while (j < Q && a[i] >= queries[j].first) {
                ans[queries[j].second] = i;
                j++;
            }
        }

        rep(i, Q) {
            cout << ans[i] + 1  << endl;
        }
    }
}

チャレンジ問題

さらに理解を深めるための問題です。↑の問題が解けた方は気が向いたらやってみましょう。

N頂点M辺の無向グラフGがある。頂点には0, 1, ..., N-1と番号がつけられており、各辺にはコストが与えられている。頂点0から頂点jまでの最短距離をd_j(ただしd_0 = 0)とする。このとき、次のQ個のクエリq_iに答えよ。

  1. $d_j \geq q_i$ かつ $|d_j - q_i|$ が最小となる $j$ (複数ある場合はそのうちで最小の $j$ )
  2. $d_j \geq q_i$ を満たす最小の $j$

ご不明な点・ご指摘等あればぜひお気軽にコメントか @yurahuna まで。

yurahuna
競プロやアルゴリズムに興味があります.
Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
No comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
ユーザーは見つかりませんでした