Total Diamonds [CodeChef December Challenge 2017 C]

https://www.codechef.com/DEC17/problems/VK18
T個の以下のクエリに答える。
　
N*Nの部屋がある。
x座標もy座標も1-indexedとすると、部屋番号はx座標とy座標の和となる。
部屋にはダイアがあり、その数は部屋番号を位毎に別々の数と考え、abs(偶数の数の総和 - 奇数の数の総和)の総和である。
全ての部屋のダイアの総和を求めよ。

解法

N=1での部屋番号を考えてみると、
1+1=2
N=2での部屋番号を考えてみると、
1+1=2
1+2=2+1=3
2+2=4
N=3での部屋番号を考えてみると、
1+1=2
1+2=2+1=3
1+3=2+2=3+1=4
2+3=3+2=5
3+3=6
の5通りとなる。
　
まず、思いつくのは部屋番号を全探索し、その部屋番号となる組合せを上手いこと求めて総和を出す解法。
これは組合せをO(1)で計算できてもO(TN)なので、ACできない。
　
上で実験してみた結果を見てみると、N=i+1の結果がN=iの結果に追加する感じになっている。
N=i+1ではN=iの結果に、1+(i+1), 2+(i+1), 3+(i+1), ... ,(i+1)+(i+1)が追加される。
(i+1)+(i+1)以外は逆の順番もあるので、同じ部屋番号が2つ追加される。
　
これを高速に計算していく。
まずは、「dia[i] := 部屋番号がiの時のダイアの数」を愚直に計算する。
次にこれを「dia[i] := 部屋番号が[0..i]の時のダイアの数の総和」と累積和にしておく。
これを使って、「ans[i] := 部屋がi*iの時のダイアの数の総和」を漸化式的に求めていく。
ans[i] = ans[i-1] + dia[(i+1)..(i*2)] + dia[(i+1)..(i*2-1)]
より、
ans[i] = ans[i-1] + dia[i*2]-dia[i] + dia[i*2-1]-dia[i]
となる。
その為クエリに答える前に答えを全部求めておき、クエリはそれから答える。

typedef long long ll;
ll dia[2010101], ans[1010101];
void pre() {
    rep(i, 0, 2010101) {
        int a = 0, b = 0, c = i;
        while (c) {
            int d = c % 10;
            if (d % 2) a += d;
            else b += d;
            c /= 10;
        }
        dia[i] = abs(a - b);
    }
    rep(i, 1, 2010101) dia[i] += dia[i - 1];

    rep(i, 1, 1000001) {
        ll a = dia[i * 2] - dia[i];
        ll b = dia[i * 2 - 1] - dia[i];
        ans[i] = ans[i - 1] + a + b;
    }
}
//---------------------------------------------------------------------------------------------------
void _main() {
    pre();

    int T; cin >> T;
    rep(t, 0, T) {
        int x; cin >> x;
        printf("%lld\n", ans[x]);
    }
}