miscalc のブログ

主に競プロの話をします

「周期性」の問題の抽象化(ライブラリ)

概要

競プロでは「周期性」を利用して解く問題がたまに出ます。例:

  • D - Teleporter

    長さ  N の数列  A = (A_1, A_2, \ldots, A_N) が与えられる( 1 \leq A_i \leq N)。 x = 1 から始めて、 x \leftarrow A_x という置き換えを  K 回行ったときの  x の値を求めよ。 N \leq 2 \times 10 ^ 5, K \leq 10 ^ {18}

  • E - Sequence Sum

    漸化式  A_1 = X, A _ {n+1} = (A _ n ^ 2 \bmod M) で定まる数列  A に対し、 \sum_{i=1}^N A_i を求めよ。 N \leq 10 ^ {10}, 0 \leq X \lt M \leq 10 ^ 5

難しくはないのですが、毎回書くと意外と頭を壊しがちなので、抽象化してライブラリにしてみました。この記事ではそのライブラリについて書きます。なお、問題の解法自体は既知として進めます。

抽象化

これらの問題は次のように抽象化できます。

モノイド  (\mathrm{Data}, \bullet) およびインデックスの集合  \mathrm{Index} があり、各インデックス  x \in \mathrm{Index} にはデータ  \mathrm{get \_ data}(x) \in \mathrm{Data} が対応づけられている。また、インデックス列  (x _ i) _ i が漸化式

 \displaystyle x _ 0  = \mathrm{start}, x _ {i+1} = \mathrm{next \_ index}(x _ i)

で定まっている。これについて、先頭からの区間

 \mathrm{get \_ data}(x _ 0) \bullet \mathrm{get \_ data}(x _ 1) \bullet \cdots \bullet \mathrm{get \_ data}(x _ {K-1})

を求めるクエリを処理したい( K がクエリで与えられる)。

これは、

  •  x _ \lambda = x _ {\lambda + \mu} となる  (\lambda, \mu) が存在する(特に最小のものをとる)。つまり、インデックスは ρ のような形で周期性を持ち、非循環部分の長さが  \lambda、循環部分の長さが  \mu である

  •  a \in \mathrm{Data}, n \in \mathbb{N} に対して冪  a ^ n = \underbrace{a \bullet \cdots \bullet a} _ {n \ \mathrm{個}} を高速に求められる(簡単のため、この冪演算およびその他必要な演算は  O(1) で行えるものとする)

という条件のもとで、前計算  O(\lambda + \mu)、クエリ  O(1) で処理できる。

「冪が高速に計算できるモノイド」である必要があるのは、( K が大きいとき)

 \begin{align}
&{\phantom{\bullet}} \left( \mathrm{get \_ data}(x _ 0) \bullet \cdots \bullet \mathrm{get \_ data}(x _ {\lambda - 1}) \right)  \tag{非循環部分} \\
&\bullet \left( \mathrm{get \_ data}(x _ \lambda) \bullet \cdots \bullet \mathrm{get \_ data}(x _ {\lambda + \mu - 1}) \right) ^ n \tag{循環部分} \\
&\bullet \left( \mathrm{get \_ data}(x _ {\lambda + n\mu}) \bullet \cdots \bullet \mathrm{get \_ data}(x _ {K - 1}) \right) \tag{あまり}
\end{align}

と変形して計算するためです*1。クエリ  O(1) にできるのは、クエリで計算する区間積が  0 スタートのものと  \lambda スタートのものしかなく前計算できるためです。

設計

コンストラク

Period<Index, Data, next_index, get_data, op, e, power> pe(start); *2

  • インデックスの型 Index演算子 != が定義されている必要があります)
  • データの型 Data
  • 一つ後のインデックスを得る関数 Index next_index(Index x)
  • インデックスからデータを得る関数 Data get_data(Index x)
  •  \mathrm{Data} 上の二項演算 Data op(Data a, Data b)
  •  (\mathrm{Data}, \bullet)単位元 Data e()
  •  a ^ n = \underbrace{a \bullet \cdots \bullet a}_{n \ \mathrm{個}} を計算する関数 Data power(Data a, long long n)

を定義する必要があります。start は最初のインデックス  x _ 0 です。

query

Data query(long long K)

 \mathrm{get \_ data}(x _ 0) \bullet \mathrm{get \_ data}(x _ 1) \bullet \cdots \bullet \mathrm{get \_ data}(x _ {K-1}) を返します。

実装

  • 既に訪れたかどうかを配列で管理するのが最も思いつきやすいと思いますが、Index の型によっては連想配列を用いる必要が出てきてしまいます。フロイドの循環検出法 *3 を用いることでこれを回避しています。

  • 限界高速化とかはしていません。

#include <iostream>
#include <vector>
using namespace std;

template <typename Index, typename Data, Index (*next_index)(Index), Data (*get_data)(Index), Data (*op)(Data, Data), Data (*e)(), Data (*power)(Data, long long)>
struct Period
{
private:
  vector<Data> dat, prod_head, prod_tail;
  int lambda, mu;
  Data head, tail;

public:
  Period(Index start)
  {
    Index a = start, b = start;
    do
    {
      dat.emplace_back(get_data(a));
      a = next_index(a);
      b = next_index(next_index(b));
    } while (a != b);
    mu = dat.size();
    Index c = start;
    while (a != c)
    {
      dat.emplace_back(get_data(a));
      a = next_index(a);
      c = next_index(c);
    }
    lambda = (int)dat.size() - mu;

    prod_head.resize((int)dat.size() + 1);
    prod_tail.resize((int)dat.size() + 1);
    prod_head[0] = e(), prod_tail[lambda] = e();
    for (int i = 1; i < (int)prod_head.size(); i++)
      prod_head[i] = op(prod_head[i - 1], dat[i - 1]);
    for (int i = lambda + 1; i < (int)prod_tail.size(); i++)
      prod_tail[i] = op(prod_tail[i - 1], dat[i - 1]);

    head = prod_head[lambda];
    tail = prod_tail[lambda + mu];
  }

  Data query(long long k)
  {
    if (k <= (int)dat.size())
      return prod_head[k];

    Data mid = prod_tail[lambda + (k - lambda) % mu];
    return op(op(head, power(tail, (k - lambda) / mu)), mid);
  }
};

使用例

 A_i をインデックス  x_i にして、 \mathrm{get \_ data}(A_i) A_i をそのまま返すようにします。二項演算  \bullet は整数の足し算  + にすればよいです。このとき power は掛け算(個数倍)になります。

using Index = long long;
using Data = long long;
int M;
Index next_index(Index x) { return x * x % M; }
Data get_data(Index x) { return x; }
Data op(Data a, Data b) { return a + b; }
Data e() { return 0; }
Data power(Data a, long long n) { return n * a; }

int main()
{
  long long N;
  int X;
  cin >> N >> X >> M;

  Period<Index, Data, next_index, get_data, op, e, power> pe(X);
  cout << pe.query(N) << endl;
}
  • D - Teleporter

    長さ  N の数列  A = (A_1, A_2, \ldots, A_N) が与えられる( 1 \leq A_i \leq N)。 x = 1 から始めて、 x \leftarrow A_x という置き換えを  K 回行ったときの  x の値を求めよ。 N \leq 2 \times 10 ^ 5, K \leq 10 ^ {18}

各操作の後の  x の値をインデックス(兼データ)にします。二項演算 op がややトリッキーです。直感的には「最も右のもの」としたいですが、このままだと単位元がないので困ってしまいます。単位元 e を「意味のないデータ」にしておいて、op は「意味のあるデータのうち、最も右のもの(なければ『意味のないデータ』)」を返すようにするとよいです。

using Index = int;
using Data = int;
vector<int> A;
Index next_index(Index x) { return A[x]; }
Data get_data(Index x) { return x; }
Data e() { return -1; }
Data op(Data a, Data b) { return b == e() ? a : b; }
Data power(Data a, long long) { return a; }

int main()
{
  int N;
  long long K;
  cin >> N >> K;
  A.resize(N);
  for (int i = 0; i < N; i++)
  {
    cin >> A[i];
    A[i]--;
  }

  Period<Index, Data, next_index, get_data, op, e, power> pe(0);
  cout << pe.query(K + 1) + 1 << endl;
}

その他の問題例

*1:単位元が必要なのは、それぞれの部分の長さが  0 になることがあるためです

*2:英語は cycle の方が正しいような気もしますが、閉路と紛らわしいので period にしました

*3:高速素因数分解のポラード・ロー法に使われていることのほうが有名?