AGC005-B：Minimum Sum

数え上げ 2変数

問題

https://beta.atcoder.jp/contests/agc005/tasks/agc005_b

N個の数 {1, 2, 3, ..., n} の任意の順列 a_n が与えられる。

\displaystyle \sum_{l=1}^N \sum_{r=l}^N min(a_l, a_{l+1}, ... ,a_r)

を求めよ。

考え方

ナイーブにシミュレーションするとO(N2)となり間に合わない。
N2 つの区間の最小値を求めて和を求めるのではなく、最小値が a_i になるときの区間がいくつ存在するか？というように見方を変えると高速に求められる。

最小値が a_i になる区間の個数を求めるために

(1)区間の左側に a_i よりも小さい数に到達するまでの距離
(2)区間の右側に a_i よりも大きな数に到達するまでの距離
を計算する。

(1) * (2) の個数が区間 [0, n-1] の中でa_i が最小値になる個数である。

サンプル3を例に考える。例えば a_i = 0 とすると、区間 [0, n-1] での左右への距離は以下のようになる。なお入力の a_i は 0-indexed としておく。

a_i 4 3 7 0 1 5 6 2
距離 3 2 1 0 1 2 3 4

左側には 3 つ、右側へは 4 つ距離を取ることができるので、合計で 1 は (3+1) * (4 + 1) = 20 回使用される。 a_i それぞれについてこの距離を求めればよい。番兵として、-1 と n = 8 をおいておくと考えやすくなる。

a_i -1 4 3 7 0 1 5 6 2 8
距離 4 3 2 1 0 1 2 3 4 5

実装上も少し工夫が必要で、入力される a_i に対して区間上の index を保持しておくと考えやすい。
以下のような形になる。

a_i 4 3 7 0 1 5 6 2
index 0 1 2 3 4 5 6 7

あとは、最小の a_i = 0 から貪欲に左右の距離を求めていく。これはTreeSetなどの平衡二分木で O(log N) で求められるので、全体の計算量は O(N log N) である。

       public void solve(int testNumber, InputReader in, PrintWriter out) {

            int n = in.nextInt();
            int[] m = new int[n];
            for (int idx = 0; idx < n; idx++) {
                int i = in.nextInt()-1;
                m[i] = idx;
            }

            TreeSet<Integer> set = new TreeSet<>();
            set.add(-1);
            set.add(n);

            long ans = 0;
            for (int i = 0; i < n; i++) {
                int idx = m[i];
                long l = idx - set.lower(idx);
                long r = set.higher(idx) - idx;
                ans += l * r * (i+1);
                set.add(idx);
            }

            out.println(ans);

        }

ポイント

  • 2変数の数え上げは片側固定でもう片方は状態をまとめて数え上げる。この問題の場合は片側固定で区間の距離を求めることがポイントであった。
  • 下から貪欲
  • indexを配列のvalueにもつ