公式

E - 家計簿と目標残高 / Household Budget and Target Balance 解説 by admin

Gemini 3.0 Flash (Thinking)

概要

この問題は、更新される数値列(入金と出金)に対して、「累積和が初めて \(K\) 以上になる位置」を効率的に求める問題です。

考察

1. 素朴な方法とその限界

各クエリに対して、1番目の取引から順番に累積和を計算していくと、1回のクエリに最大で \(O(N)\) の時間がかかります。クエリの数を \(Q\) とすると、全体の計算量は \(O(NQ)\) となり、制約の \(N, Q \leq 2 \times 10^5\) では実行時間制限に間に合いません。

2. 累積和の性質と更新

取引が更新されるため、累積和をあらかじめ計算しておく手法もそのままでは使えません。また、出金(負の値)があるため、累積和は単調増加とは限りません。そのため、累積和の配列に対して二分探索を適用することも困難です。

3. セグメント木による高速化

「区間の合計値」と「区間内での最大累積和」を管理できるセグメント木を利用することで、更新と探索の両方を高速化できます。 各ノードに以下の2つの情報を持たせます。 - sum: その区間に含まれる取引金額の合計 - max_p: その区間の左端から始めたときの、区間内での累積和の最大値

2つのノード(左の子 L、右の子 R)を結合して親ノードを作る際、親の max_p は以下のいずれか大きい方になります。 1. 左の子の区間内で最大値に達する場合:L.max_p 2. 左の子をすべて通り過ぎ、右の子の区間内で最大値に達する場合:L.sum + R.max_p

このように情報を保持することで、セグメント木の上を二分探索のように降りていく「セグメント木上の二分探索」が可能になります。

アルゴリズム

セグメント木の構築

  1. 各葉ノードに取引の値をセットします(R なら正、L なら負)。
  2. 下から上へ、上述のルール(summax_p)に従ってノードを埋めます。

タイプ 1:取引の修正

指定されたインデックスの葉ノードを更新し、ルート(根)に向かって summax_p を再計算します。これは \(O(\log N)\) で可能です。

タイプ 2:目標残高の探索

ルートから順に以下の判定を行い、目的の取引番号を探します。 - 現在の累積和を current_sum とします。 - 左の子の max_p を使い、current_sum + left.max_p >= K かどうかを判定します。 - YES の場合: 左の子の区間のどこかで初めて \(K\) を超えることが確定するため、左の子へ進みます。 - NO の場合: 左の子の区間内では \(K\) に達しないため、右の子へ進みます。このとき、current_sum に左の子の sum を加算しておきます。

この探索も、木の高さ分(\(O(\log N)\))のステップで完了します。

計算量

  • 時間計算量: \(O((N + Q) \log N)\)
    • セグメント木の構築に \(O(N)\)、各クエリ(更新・探索)に \(O(\log N)\) かかります。
  • 空間計算量: \(O(N)\)
    • セグメント木のノード数は最大で \(4N\) 程度必要です。

実装のポイント

  • 大きな値の扱い: 累積和は \(2 \times 10^{14}\) 程度になる可能性があるため、C++では long long 型を使用する必要があります。

  • 探索の終了条件: ルートノードの max_p\(K\) 未満であれば、配列のどこを探索しても \(K\) に達しないため、即座に 0 を出力できます。

  • セグメント木のサイズ: \(N\) を超える最小の 2 のべき乗を葉の数とすることで、実装がシンプルになります。

    ソースコード

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>

using namespace std;

/**
 * Problem: Household Account Book and Target Balance
 * We need to maintain prefix sums of transactions and find the first transaction 
 * such that the prefix sum reaches or exceeds a given target K.
 * Transactions can be additions (R) or subtractions (L).
 *
 * Approach:
 * Use a Segment Tree where each node stores:
 * 1. The total sum of elements in its range.
 * 2. The maximum prefix sum within its range (relative to the start of the range).
 *
 * Time Complexity:
 * - Build: O(N)
 * - Update: O(log N)
 * - Query (Type 2): O(log N)
 * Total: O((N + Q) log N)
 */

struct Node {
    long long sum;
    long long max_p;
};

// N <= 2e5, so tree_size is up to 2^18 = 262144.
// 2 * tree_size is up to 524288. MAX_NODES = 1 << 19 is sufficient.
const int MAX_NODES = 1 << 19;
Node tree[MAX_NODES];
int tree_size;

/**
 * Update the p-th transaction (0-indexed) with a new value.
 */
void update(int i, long long val) {
    i += tree_size;
    tree[i] = {val, val};
    while (i > 1) {
        i /= 2;
        int left = 2 * i;
        int right = 2 * i + 1;
        tree[i].sum = tree[left].sum + tree[right].sum;
        // The max prefix sum of the merged range is the maximum of:
        // 1. Max prefix sum of the left child.
        // 2. Sum of the left child + max prefix sum of the right child.
        tree[i].max_p = max(tree[left].max_p, tree[left].sum + tree[right].max_p);
    }
}

/**
 * Traverse the segment tree to find the first index where the global prefix sum >= K.
 * current_sum is the sum of all elements to the left of the current node's range.
 */
int find_first(int node, int l, int r, long long current_sum, long long K) {
    if (l == r) {
        return l;
    }
    int mid = l + (r - l) / 2;
    int left_child = 2 * node;
    int right_child = 2 * node + 1;

    // Check if the target K is reached within the left child's range.
    if (current_sum + tree[left_child].max_p >= K) {
        return find_first(left_child, l, mid, current_sum, K);
    } else {
        // Otherwise, the first occurrence must be in the right child's range.
        return find_first(right_child, mid + 1, r, current_sum + tree[left_child].sum, K);
    }
}

int main() {
    // Fast I/O
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    if (!(cin >> N >> Q)) return 0;

    // Determine the smallest power of 2 >= N for the segment tree leaves.
    tree_size = 1;
    while (tree_size < N) tree_size *= 2;

    // Initialize all nodes in the tree.
    // max_p is initialized to a very small value so that "empty" leaves beyond N
    // do not affect the prefix sum calculations.
    for (int i = 0; i < MAX_NODES; ++i) {
        tree[i] = {0, -4000000000000000LL};
    }

    // Read initial transactions and place them in the segment tree leaves.
    for (int i = 0; i < N; ++i) {
        char d;
        long long b;
        cin >> d >> b;
        long long val = (d == 'R' ? b : -b);
        tree[tree_size + i] = {val, val};
    }

    // Build the segment tree from bottom to top.
    for (int i = tree_size - 1; i >= 1; --i) {
        int left = 2 * i;
        int right = 2 * i + 1;
        tree[i].sum = tree[left].sum + tree[right].sum;
        tree[i].max_p = max(tree[left].max_p, tree[left].sum + tree[right].max_p);
    }

    // Process Q queries.
    for (int q_idx = 0; q_idx < Q; ++q_idx) {
        int type;
        cin >> type;
        if (type == 1) {
            // Point update: 1 p D b
            int p;
            char d;
            long long b;
            cin >> p >> d >> b;
            long long val = (d == 'R' ? b : -b);
            update(p - 1, val);
        } else if (type == 2) {
            // Search query: 2 K
            long long K;
            cin >> K;
            // tree[1].max_p is the maximum prefix sum across the entire array.
            if (tree[1].max_p < K) {
                cout << 0 << "\n";
            } else {
                // find_first returns 0-indexed position; convert to 1-indexed for output.
                cout << find_first(1, 0, tree_size - 1, 0, K) + 1 << "\n";
            }
        }
    }

    return 0;
}

この解説は gemini-3-flash-thinking によって生成されました。

投稿日時:
最終更新: