集合内の任意の2要素の差の二乗の総和 - 素数好きの最高技術責任者のブログ

先日、Facebook社が年1回開催しているプログラミングコンテスト（以下、プロコン）である HackerCup というものに参加しました。自分はプロコンガチ勢ではなく趣味でたまに参加する程度なのですが、恥ずかしながらこのHackerCupの存在を今まで知らず、今年初めて参加してみたのでした。

QualificationRound（予選）、Round1，Round2、Round3、FinalRoundと続く中、自分は残念ながらRound2で敗退してしまったのですが、そのRound2で出題された問題の中で使える面白いテクニックを知るにあたりそこはかとなくテンションがあがったので、紹介したいと思います。プロコンガチ勢であれば当たり前のように使っているであろうテクニックらしいですが、自分は今回初めて知ったので、少しでも多くの人にこのテクニックが伝わればいいな、と。またプロコンに興味が無い人が少しでも興味を持ってくれればいいな、と。あるいは世の中に存在する遅いコードが一部でも高速になればいいな、と。そんな思いで書くことにします。（っていうか、珠玉のプログラミングとかハッカーのたのしみとかプログラミングコンテストチャレンジブックとかに書いてありそうな話題ですけど、気にしないことにします。）

では早速ですが問題です。

N個の整数、
$\displaystyle X_1, X_2, X_3, ... X_N$
があります。
この中から任意の2つを選ぶ方法は
$\displaystyle \begin{eqnarray}{}_N C _2\end{eqnarray} = \frac{N * (N - 1)}{2}$
通りありますが、その全ての選び方で選んだ2つの整数の差の2乗の総和を求めなさい。ただし、総和は相当大きな数値になるので、総和を $1000007(=10^6 + 7)$ で割った余りを求めなさい。

制限として、各値は
$\displaystyle 2 \leq N \leq 1,000,000$
$\displaystyle -100,000,000 \leq X_i \leq 100,000,000 (1 \leq i \leq N)$
を満たすものとします。

なお、入力ファイルは、1行目にNの値、2行目からN+1行目にX₁からX_Nが1行に1つずつ記されているものとします。

入力ファイルサンプル1

サンプル1の正解は、 $\displaystyle (-1 - 0)^2 + (-1 - 1)^2 + (0 - 1)^2$ = 6となります。

入力サンプル2

サンプル2の正解は 181270 です。

さて、この時この入力ファイル（サイズ約9MB）に対して正解を求められますか？というのが今回の問題になります。N=1,000,000です。

この問題を見て「あぁアレね」とピンと来ていないエンジニアの方は、是非考えてみてください。多分楽しいので。もしあなたがエンジニアではなければ、身近にいるエンジニアに出題してみるのも面白いかもしれません。

以下、解説です。

まず、最もシンプルで直感的なコードは以下のような感じになると思います。

def solve(N, X, MOD=10 ** 6 + 7):
    """
        >>> solve(3, [-1, 0, 1])
        6
        >>> solve(5, [293842, 238424, 174234, 924234, 532800])
        181270
    """
    assert N == len(X)

    total = 0
    for i in xrange(N):
        for j in xrange(i + 1, N):
            total += (X[i] - X[j]) ** 2
            total %= MOD
    return total

何をやっているかというと、2重のforループ（ $O(N^2)$ ）で全組み合わせを計算しています。

これだとサンプルの問題は解けても、おそらくNが5,000を超えたあたりから「なんか時間かかるな…」という感じになり、上述の入力ファイルに対してはPCを丸1日ぶん回せば解けるかも？…という感じになるかと思います。要するに遅いプログラムなわけです。

さて、これを数秒以内に解く方法が以下になります。

まず、正解を求める関数をf(N)とします。（実際はNの関数ではなくN個の要素を引数とする関数ですが、わかりやすさのために。）

そしてf(N)が計算されているとして、f(N+1)を考えると、

$\displaystyle f(N+1) = f(N) + \sum_{i=1}^{N}(X_{N+1}-X_i)^2$

であることがすぐにわかります。

これを変形すると、

$\begin{align} f(N+1) - f(N) &= \sum_{i=1}^{N}(X_{N+1}^2-2X_{N+1}X_i+X_i^2) \\\ &= NX_{N+1}^2 - 2X_{N+1}\sum_{i=1}^{N}X_i + \sum_{i=1}^{N}X_i^2 \end{align}$

となります。

つまり、X_iの総和とX_i²の総和がわかっていれば、いちいちX_N+1と各要素X_iの差を計算しなくても済むのです。

というわけで、先のプログラムを以下のようにfor文の1重ループ（ $O(N)$ ）に書き換えることができます。

def solve(N, X, MOD=10 ** 6 + 7):
    """
        >>> solve(3, [-1, 0, 1])
        6
        >>> solve(5, [293842, 238424, 174234, 924234, 532800])
        181270
    """
    assert N == len(X)

    total = 0
    sum_x = sum_x2 = 0
    for i in xrange(N):
        total += i * X[i] ** 2 - 2 * X[i] * sum_x + sum_x2
        total %= MOD
        sum_x += X[i]
        sum_x2 += X[i] ** 2
    return int(total)

このテクニックは2乗じゃなくて3乗になったりk乗になっても使えるし、「集合内の任意の2要素」じゃなくて3要素にしても使えそうだし、k乗じゃなくてbit演算とかが混じっても面白そうだし、シンプルなゆえに色々応用が利きそうで楽しいですね。