D - Kadomatsu Subsequence 解説 by MMNMM


\(\max\lbrace i,j,k\rbrace=j\) であるようなものをすべて数え上げることができれば、列を反転させてもう一度解き、答えを足せばよいです。 そのような数え上げについて考えます。

数え上げるべきものは、次の条件を満たす整数の \(3\) つ組 \(i,j,k\) です。

  • \(i\lt j\) かつ \(k\lt j\)
  • \(\dfrac{A _ i}3=\dfrac{A _ k}7=\dfrac{A _ j}5\)

正の整数で番号付けられた列 \(T=(T _ 1,T _ 2,\ldots),S=(S _ 1,S _ 2,\ldots)\) を用意し、はじめすべて \(0\) にしておきます。 \(i\) の昇順に \(A _ i\) を見て、次の操作を行います。

  • \(A _ i\) が \(3\) の倍数のとき、\(T _ {A _ i/3}\) を \(1\) 増やす。
  • \(A _ i\) が \(7\) の倍数のとき、\(S _ {A _ i/7}\) を \(1\) 増やす。
  • \(A _ i\) が \(5\) の倍数のとき、答えに \(T _ {A _ i/5}S _ {A _ i/5}\) を加える。

列 \(T,S\) を連想配列などを用いて管理することで、この問題を \(O(N\log N)\) 時間や \(O(N)\) 時間で解くことができます。

実装例は以下のようになります。

#include <iostream>
#include <vector>
#include <map>
#include <ranges>

int main() {
    using namespace std;
    // i < j, k < j を数え上げる関数
    const auto solve{[](const auto& seq) {
        map<unsigned, unsigned> three, seven;
        unsigned long ans{};
        for (const auto x : seq) {
            if (x % 3 == 0)
                ++three[x / 3];
            if (x % 7 == 0)
                ++seven[x / 7];
            if (x % 5 == 0)
                ans += three.contains(x / 5) && seven.contains(x / 5) ? static_cast<unsigned long>(three[x / 5]) * seven[x / 5] : 0;
        }
        return ans;
    }};
    const auto A{views::istream<unsigned>(cin) | views::drop(1) | ranges::to<vector<unsigned>>()};
    // 順方向と逆方向を解き、足し合わせる
    cout << solve(A) + solve(A | views::reverse) << endl;
    return 0;
}

投稿日時:
最終更新: