裏紙

ほぼ競プロ、たまに日記

WUPC 2019 H - Doki Doki Programming Clubs!

問題

H - Doki Doki Programming Clubs!

問題概要

N個の数列が与えられる。 i番目の数列A_iの長さはK_iである。 クエリ(X,Y)とは以下の通り:

  • 数列A_X, A_Yについて、 \displaystyle \sum_{i=0}^{\max (K_X, K_Y)-1} A_{X, \ i \  \% \  K_X} \times A_{Y,\ i \  \% \  K_Y}10^9 + 7で割った余りを答える

上記のクエリがQ個与えられるので、各クエリについて答えよ。

  •  2 \le N \le 200000
  •  1 \le K_i
  •  \sum K_i \le 200000
  •  1 \le A_{i,j} \le 10^9
  •  1 \le Q \le 200000
  •  0 \le X_i , Y_i \le N-1
  •  X_i \neq Y_i

イデア

まず、クエリについて、どのような処理になっているかを観察してみる。 以下では、クエリ(x,y)について、K_x \le K_yを仮定する。 数列A_yの方が長いので、短い方の数列A_xを長い方に揃えることになる。 例えば、K_x = 3, K_y = 8ならこんな感じ。

f:id:imulan:20190312214500p:plain

短い方の数列A_xを長い方に周期的に揃えて、積の和について考える。

f:id:imulan:20190312214716p:plain

ここで、問題で与えられた通りに短い方を長い方に毎回揃えて計算してしまうとTLEするので、短い方に揃えて計算出来ないかを考える。 長い方のindexに注目してみると、短い方の長さK_xの剰余によって、どの要素との積になっているかが決定している:

f:id:imulan:20190312215153p:plain

よって、クエリに対して O( \min (K_x, K_y))で計算したい時、数列A_yを周期K_xで分類した和(上の例で言うと、数列A_yに関する赤の和と青の和と緑の和)を用意しておけば良いことになる。

以上の考察から、短い方の数列に揃えてクエリを処理するためには、前計算として各数列A_iについて、自分より短い周期について、その剰余で分類した和を用意しておく必要がある。

以下では、 \sum K_i = Sとする。 この前計算は O(S \sqrt{S})で実現することができる。 その理由については、実際に前計算のコードを見てみるとわかりやすい:

// 前計算
// k[i] のdistinctな要素を列挙
vector<int> uk(k);
sort(uk.begin(), uk.end());
uk.erase(unique(uk.begin(), uk.end()), uk.end());
int UK = uk.size();

// 各数列 a[i] について、自分以下の周期の和を前計算
vector<vector<vector<long>>> sum(n, vector<vector<long>>(UK));
for(int i=0; i<n; ++i){
    for(int j=0; j<UK; ++j){
        int L = uk[j];
        if(L > k[i]) break;

        sum[i][j] = vector<long>(L);
        for(int l=0; l<k[i]; ++l){
            sum[i][j][l%L] += a[i][l];
            sum[i][j][l%L] %= mod;
        }
    }
}

iのループとlのループを合わせて、ちょうどS回のループになっているので、ループが回る回数はO(S \times UK)となっている。 ここで、UKとは、K_iの種類数となっているが、数列の長さK_iの種類数はたかだかO( \sqrt{S})である。

できるだけ種類数が多くなるような状況を考えてみると、数列の長さは1,2,3,4, \ldotsと、長さを1ずつ増やしながら1つずつ数列を用意する、という状況になる。 これを長さtまでもれなく揃えたとすると、長さの合計は1+2+ \ldots + t = \frac{t(t+1)}{2}となり、\frac{t(t+1)}{2} \le Sなので、種類数はO( \sqrt{S})で抑えられていることが分かる。

以上で、前計算は O(S \sqrt{S})であることが確認できた。 以下ではクエリ処理について確認する。

一度来たクエリについてはメモ化して答えることにし、新しいクエリについては O( \min (K_x, K_y)) で計算することにすると、計算量は O( (S+Q) \sqrt{S} ) となっている。 これを示すために、2つの場合に分けて考える:

 \min (K_x, K_y) \lt \sqrt{S}のとき

1クエリに対して O( \min (K_x, K_y)) = O(\sqrt{S})で処理できるので、クエリ全体についてO(Q \sqrt{S})で処理できる。

 \min (K_x, K_y) \ge \sqrt{S}のとき

この場合は、K_i \ge \sqrt{S}となるような数列はたかだか\sqrt{S}個しか存在しない。 この条件を満たす数列A_iに対して、i以外の数列とのクエリが全部与えられたとしても、そのクエリ合計でO(S)で抑えられているのは明らか。 以上のことから、クエリ全体についてO(S \sqrt{S})で処理できる。

以上より、計算量は全体で O((S+Q) \sqrt{S})となり、これは十分に高速である。

実装(C++)

#include <bits/stdc++.h>
using namespace std;

const long mod = 1e9+7;

int main(){
    int n;
    scanf(" %d", &n);

    vector<int> k(n);
    vector<vector<long>> a(n);
    for(int i=0; i<n; ++i){
        scanf(" %d", &k[i]);
        a[i] = vector<long>(k[i]);
        for(int j=0; j<k[i]; ++j) scanf(" %ld", &a[i][j]);
    }

    // 前計算
    // k[i] のdistinctな要素を列挙
    vector<int> uk(k);
    sort(uk.begin(), uk.end());
    uk.erase(unique(uk.begin(), uk.end()), uk.end());
    int UK = uk.size();

    // 各数列 a[i] について、自分以下の周期の和を前計算
    vector<vector<vector<long>>> sum(n, vector<vector<long>>(UK));
    for(int i=0; i<n; ++i){
        for(int j=0; j<UK; ++j){
            int L = uk[j];
            if(L > k[i]) break;

            sum[i][j] = vector<long>(L);
            for(int l=0; l<k[i]; ++l){
                sum[i][j][l%L] += a[i][l];
                sum[i][j][l%L] %= mod;
            }
        }
    }

    // クエリ処理
    int Q;
    scanf(" %d", &Q);

    map<pair<int,int>, long> memo;
    while(Q--){
        int x,y;
        scanf(" %d %d", &x, &y);

        pair<int,int> p(min(x,y), max(x,y));
        long ans = 0;
        if(memo.count(p)) ans = memo[p];
        else{
            if(k[x] > k[y]) swap(x,y);
            int uk_idx = lower_bound(uk.begin(), uk.end(), k[x]) - uk.begin();
            for(int i=0; i<k[x]; ++i){
                ans += a[x][i] * sum[y][uk_idx][i];
                ans %= mod;
            }
            memo[p] = ans;
        }
        printf("%ld\n", ans);
    }
    return 0;
}