yukicoder No.956 Number of Unbalanced

https://yukicoder.me/problems/no/956

問題概要

N要素の正数列Aが与えられる。Aのすべての連続する部分列のうち、ある種類の要素が部分列の過半数を占めているものが何通りあるか求めよ。

N <= 105

解法

想定解はO(NlogN)だったが、O(N{\sqrt{N}})で解いたので書き残しておく。このやり方では数列中に出現する数を「出現回数が{\sqrt{N}}以上か未満か」で2つに分けたうえで、それぞれを独立に解く。基本的な方針は以下のようになる。

  • (1) 出現回数が{\sqrt{N}}未満の場合、その数が中央値になり得る区間長はたかだか2*{\sqrt{N}}なので、そこまでの各区間長について個別にO(N)で処理する。
  • (2) 出現回数が{\sqrt{N}}以上の場合、そのような数はたかだか{\sqrt{N}}種類までしかありえないので、それぞれの数をO(N)で処理する。

(1) 出現回数が{\sqrt{N}}未満の数の場合

この場合は条件を満たす数をまとめて処理していくことにする。具体的には上の通り区間長1〜2*{\sqrt{N}}の場合についてそれぞれ個別にアンバランスになる区間の数をかぞえていく。これはローリングハッシュとかでよくやるように区間をひとつずつずらしながら「現在見ている区間における各数の出現回数」を管理していけばよい。管理のためには一個の配列を持っていれば十分で、区間をずらしたときに外れた数を-1, 入ってきた数を+1するだけ。そしてこの足し引きの結果ちょうど過半数に到達したか、あるいは逆にちょうど過半数を下回ったかを見ていれば「現在見ている区間がアンバランスか」は容易に判定できる。

区間長がO({\sqrt{N}})個で、各区間長をO(N)で処理できるので、全体の計算量はO(N{\sqrt{N}})となる。

(2) 出現回数が{\sqrt{N}}以上の数の場合

この条件に当てはまる数はたかだか{\sqrt{N}}個しかないため、それぞれの数について個別に考えていくことにする。解くべき問題は「xかそれ以外かで埋まった数列が与えられる。xが過半数を占める区間は何通りか」である。

ひとつのやり方として、これはxを1, それ以外を-1とみなして累積和をとることで効率的に求めることができる。この変換によって ある区間でxが過半数区間和が正 となるので、累積和テーブルを左から見ていって「過去の値の中で今見ている値より小さいものが何個あるか」を逐一数え上げていけばよい。これは素朴にやろうとするとsetとかでlogがついてしまうが、ここで考えている数列の「+1か-1しかない」という性質を使うと各インデックスにつきO(1)で済ますことができる。要は「累積和テーブルにおいて今見ている値より小さい値の出現回数」がわかっていればいいので、「今見ている値」「今見ている値より小さい値の出現回数」「それぞれの値が何回出現したかをすべて覚えておくテーブル」を用意しておいたうえで、累積和テーブルを1進んだときにそれぞれをよしなに足し引きしてやればよい(+1の場合「今見ている値」は+1, 「今見ている値より小さい値の出現回数」は+table[今見ている値-1], テーブルはtable[今見ている値] += 1. いう感じ)。あとはこの処理のあと毎回「今見ている値より小さい値の出現回数」を答えに足し込んでいけば答えが求まる。

数の種類が{\sqrt{N}}, 各数の処理がO(N) なのでこれもO(N{\sqrt{N}})となり、(1)(2)あわせて全体としても結局O(N{\sqrt{N}})となる。だいたい配列を舐めながら一時変数を弄っているだけなのでC++なら定数倍も軽く済む。

感想

sqrt(N)を機械的にマークダウンのtexでくるんだらキモい表示になってしまった

想定解まだちゃんとわかってないけど高速化のために使っている性質は同じっぽい?

コード (C++)

#include <bits/stdc++.h>
using namespace std;
#define REP(i,n) for (int i=0;i<(n);i++)
#define REP2(i,m,n) for (int i=m;i<(n);i++)
typedef long long ll;

const int SQ = 320;
const int OFFSET = 101010;

int N;
int A[101010];
int B[101010];
int C[101010];
ll D[202020];
int E[101010];
int cnt[101010];

void solve() {
    cin >> N;
    REP(i, N) cin >> A[i];
    REP(i, N) cnt[A[i]] += 1;

    vector<int> less, much;
    REP(i, 101010) if (cnt[i] > 0) (cnt[i] < SQ ? less : much).push_back(i);
    REP(i, N) (cnt[A[i]] < SQ ? B[i] : C[i]) = A[i];

    int L = less.size();
    int M = much.size();
    ll ans = 0;

    REP2(len, 1, min(N+1, SQ*2)) {
        for (auto a : less) {
            E[a] = 0;
        }
        bool ok = false;
        const int half = len / 2 + 1;
        REP(i, len) if (B[i] != 0) {
            E[B[i]] += 1;
            if (E[B[i]] == half) ok = true;
        }
        ans += ok;
        REP2(i, 1, N-len+1) {
            if (i > 0 && B[i-1] != 0) {
                if (E[B[i-1]] == half) ok = false;
                E[B[i-1]] -= 1;
            }
            if (B[i+len-1] != 0) {
                E[B[i+len-1]] += 1;
                if (E[B[i+len-1]] == half) ok = true;
            }
            ans += ok;
        }
    }

    for (const auto a : much) {
        memset(D, 0, sizeof(D));
        int pointer = OFFSET;
        ll counter = 0; // pointer 未満の合計
        D[OFFSET] += 1;
        REP(i, N) {
            if (A[i] == a) {
                pointer += 1;
                counter += D[pointer - 1];
            } else {
                counter -= D[pointer - 1];
                pointer -= 1;
            }
            D[pointer] += 1;
            ans += counter;
        }
        //cout << pointer - OFFSET << " " << counter << " " << D[pointer] << endl;
    }

    cout << ans << endl;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    solve();
}

yukicoder No.951 【本日限定】1枚頼むともう1枚無料!

https://yukicoder.me/problems/no/951

問題概要

N枚のピザがあり、ピザには値段Pi円と美味しさDiがそれぞれ設定されている。またピザを1枚買うたび、そのピザ以下の値段のピザを1枚無料で買うことができる。K円まで使えるとき、買ったピザの美味しさの総和としてありえる最大値を求めよ。

N, K, Di <= 5000

Pi <= K

解法

いかにもDPでナップサックしてくださいという見た目だが、単純なやり方ではタダ券の処理が難しい。たとえば高い順に見ていってタダ券が何個余っているかをDPのキーに含めようとすると、キーが【何番目まで見た, 何円使った, 何個余ってる】で最悪50003になってしまう。しかしここで「買い物が終わったあとの最終的な形」について思いを馳せると、買う商品を固定した場合、高い順に「買う→タダにする→買う→タダにする」を交互にやるのがもっとも安く済む方法である。というわけで実は高い順に見ていったとき「タダ券を持っているなら次の買い物で即使う」のが最適であるので、DPのキーには「タダ券がある・ない」の2値を含めるだけで十分である。

DP[i][j][k] = (あらかじめ商品を値段の高い順にソートしたとき)i番目の商品まで見て、j円使っていて、タダ券がk枚(k=0, 1)あるときの得られる美味しさの総和の最大値

このDPは状態数がO(NK)で遷移は買う・買わないの2通りのO(1)となるので計算量はO(NK).

感想

雑な一般論だが、DPしたいけど状態が多すぎるというとき実はある一部分が貪欲に処理できるので状態数を減らせるというパターンは結構あって、この問題のような「最終形を想像すると最適な操作や並び順が固定できる」というのもその類型のひとつというような印象がある(この間のABCでも似たようなやつがあった)

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;

void main() {
    auto s = readln.split.map!(to!int);
    auto N = s[0];
    auto M = s[1];
    auto A = iota(N).map!(_ => readln.split.map!(to!int).array).array;
    A.sort!"a[0] > b[0]";

    auto dp = new int[][][](2, M+1, 2);
    foreach (i; 0..2) foreach (j; 0..M+1) dp[i][j][1] = -(1 << 29);

    int cur = 0;
    int tar = 1;

    foreach (i; 0..N) {
        int p = A[i][0];
        int d = A[i][1];
        foreach (j; 0..M+1) {
            dp[tar][j][0] = 0;
            dp[tar][j][1] = -(1 << 29);
        }
        foreach (j; 0..M+1) {
            dp[tar][j][0] = max(dp[tar][j][0], dp[cur][j][0]);
            dp[tar][j][1] = max(dp[tar][j][1], dp[cur][j][1]);
            if (p + j <= M) dp[tar][j+p][1] = max(dp[tar][j+p][1], dp[cur][j][0] + d);
            dp[tar][j][0] = max(dp[tar][j][0], dp[cur][j][1] + d);
        }
        swap(cur, tar);
    }

    int ans = 0;
    foreach (j; 0..M+1) foreach (k; 0..2) ans = max(ans, dp[cur][j][k]);
    ans.writeln;
}

Codeforces Round #344 (Div. 2): D. Messenger

https://codeforces.com/problemset/problem/631/D

問題概要

文字列S, Tが与えられる。SがTを部分文字列として何個含むか求めよ。ただしS, Tは(回数, 文字)というタプルの列として与えられる。これはその文字が何回連続するかを表した形式で、例えば(3, a)(2, b)(4, c)はaaabbccccという文字列を表す。

S, Tに対応するタプル列の長さ <= 2×105

各タプル中の「回数」の値 <= 106

解法

まず表記の形式をユニークにするため、同じ文字に対応するタプルが連続しているところはひとつにまとめてしまう((3, a)(4, a)(1, a)のような部分は(8, a)にまとめる)。このようにすると、TがSの部分文字列として現れるならば対応箇所のタプル列の長さがSとTで等しくなる。これによって以下のような場合分けができる。

Tのタプル列の長さが1の場合

これは単純にSの要素をひとつずつ見ていって愚直に数えれば良い。

Tのタプル列の長さが2の場合

これもSの要素を2つずつ見ていけば数えられる。

Tのタプル列の長さが3以上の場合

このときSのある箇所がTの部分文字列であるための条件は

  • すべてのタプルについて、文字が一致する
  • 両端以外のタプルについて、回数が一致する
  • 両端のタプルについて、Sの回数がTの回数以上である

のすべてが満たされることである。そしてこのような部分をSから見つけ出すためには

  • 文字だけに着目して文字列検索を行い、Sの中でTと一致する箇所を列挙する
  • (両端以外の)回数だけに注目して「数字列」の検索を行い、Sの中でTと一致する箇所を列挙する

以上2つの前処理を行っておけばよい。前者で列挙した文字列一致インデックスを舐めていき、後者で列挙した数字列インデックスに対応する箇所(両端以外の列)が存在しているか、またS側の両端の回数はTの両端の回数と一致しているかを確認する。すべて満たされていればカウント+1となる。

肝心のこの前処理をどう効率的に行うかだが、これは単純な文字列検索なのでKMP法によってO(len(S) + len(T))でできる。

感想

KMP法を初めて書いたので記念カキコ(wikipedia擬似コードを丸写ししただけ)

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.stdio;

void main() {
    auto s = readln.split.map!(to!int);
    long[] A_INT, B_INT;
    string A_STR, B_STR;

    void parse(ref long[] x, ref string y) {
        auto t = readln.split;
        foreach (a; t) {
            auto v = a.split("-").front.to!long;
            auto c = a.split("-").back[0];
            if (y.empty || y.back != c) {
                y ~= c;
                x ~= v;
            } else {
                x.back += v;
            }
        }
    }

    parse(A_INT, A_STR);
    parse(B_INT, B_STR);
    int N = A_INT.length.to!int;
    int M = B_INT.length.to!int;
    long ans = 0;

    if (M == 1) {
        foreach (i; 0..N) {
            if (A_STR[i] != B_STR.front) continue;
            ans += max(0L, A_INT[i] - B_INT.front + 1);
        }
    } else if (M == 2) {
        foreach (i; 0..N-1) {
            if (A_STR[i..i+2] != B_STR) continue;
            if (A_INT[i] >= B_INT[0] && A_INT[i+1] >= B_INT[1]) ans += 1;
        }
    } else {
        auto int_idx = kmp_search(A_INT, B_INT[1..M-1]);
        auto str_idx = kmp_search(A_STR, B_STR);
        bool[int] dict;
        foreach (idx; int_idx) dict[idx-1] = true;
        foreach (idx; str_idx) {
            if (A_INT[idx] >= B_INT[0] && A_INT[idx+M-1] >= B_INT[M-1] && idx in dict) {
                ans += 1;
            }
        }
    }

    ans.writeln;
}

int[] kmp_search(T)(const T[] S, const T[] W) {
    int j, k;
    int[] P;
    auto table = kmp_table(W);

    while (j < S.length) {
        if (W[k] == S[j]) {
            j += 1;
            k += 1;
            if (k == W.length) {
                P ~= j - k;
                k = table[k];
            }
        } else {
            k = table[k];
            if (k < 0) {
                j += 1;
                k += 1;
            }
        }
    }

    return P;
}

int[] kmp_table(T)(const T[] W) {
    auto n = W.length.to!int;
    auto table = new int[](n+1);
    table[0] = -1;
    int pos = 1, cnd = 0;

    while (pos < n) {
        if (W[pos] == W[cnd]) {
            table[pos] = table[cnd];
        } else {
            table[pos] = cnd;
            cnd = table[cnd];
            while (cnd >= 0 && W[pos] != W[cnd]) {
                cnd = table[cnd];
            }
        }
        pos += 1;
        cnd += 1;
    }

    table[pos] = cnd;
    return table;
}

Codeforces Round #259 (Div. 1): B. Little Pony and Harmony Chest

https://codeforces.com/problemset/problem/453/B

問題概要

N要素の正整数列Aが与えられる。すべての相異なる2要素が互いに素であるようなN要素の正整数列であって、Σ(abs(A[i]-B[i])) が最小となるような数列Bを構成せよ。

N <= 100

Ai <= 30

解法

相異なる2要素が互いに素であるという条件から「1は何回でも使える」「1以外の数は最大でも1回しか使えない」ことがいえる。またAiが最大でも30までという制約の小ささに注目すると、1が何回でも使えることからBに対して59以上の数は使う意味がないということがわかる(代わりに1を使えば必ずAとの差は小さくなるので)。さらに最小性の条件から、Bにおける各要素の大小関係はAと同じとしてよい(Ai > Aj かつ Bi < Bj だった場合、BiとBjをswapすれば必ず最終スコアは小さくなる)。また互いに素の条件から、ある数を使ったときそれとのGCDが1でない数は以降使えなくなる。使える数はたかだか60なのでどの数がもう使えないかは64bit型整数で管理することができる。

以上の条件をすべて踏まえると、【今見ているインデックス、使える数の最大値、もう使えない数のビットマスク、現在の差の合計】あたりの値をもってDFSで全探索していくことができ、制限時間も長いのでこれで十分間に合う。

感想

計算量がまったくわからない

参考にした記事

Codeforces #259 Div1 B. Little Pony and Harmony Chest - kmjp's blog

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.stdio;

void main() {
    auto s = readln.split.map!(to!int);
    auto n = s[0];
    auto A = readln.split.map!(to!int).array;

    auto B = n.iota.map!(i => tuple(A[i], i)).array;
    B.sort!"a[0] > b[0]";

    auto C = new int[](n + 1);
    foreach (i; 0..n) C[i + 1] += C[i] + A[i];

    auto mask = new long[](60);
    foreach (i; 2..60) foreach (j; 2..60) if (gcd(i, j) != 1) mask[i] |= 1L << j;

    int[] ans;
    int[] tmp = new int[](n);
    tmp[] = 1;
    int minv = 1 << 29;

    void dfs(int idx, int num, long used, int val) {
        if (idx == n) {
            if (val < minv) {
                minv = val;
                ans = tmp.dup;
            }
            return;
        }
        if (num == 1) {
            int fval = val + C[n] - C[idx] - (n - idx);
            if (fval < minv) {
                minv = fval;
                ans = tmp.dup;
            }
            return;
        }
        foreach_reverse (v; 1..num+1) {
            if (used & (1L << v)) continue;
            tmp[B[idx][1]] = v;
            dfs(idx+1, max(1, v-1), used | mask[v], val + abs(B[idx][0] - v));
            tmp[B[idx][1]] = 1;
        }
    }

    dfs(0, 58, 0, 0);
    ans.map!(to!string).join(" ").writeln;
}

AtCoder Regular Contest 090: F - Number of Digits

https://atcoder.jp/contests/arc090/tasks/arc090_d

問題概要

正整数nに対してf(n)をnの10進表記での桁数と定義する。与えられる整数Sに対してf(l)+f(l+1)+...+f(r)=Sを満たすような(l, r)の組が何通り存在するか求めよ。

S <= 108

解法

まず桁数が少しでも大きくなると「ある桁全体を覆うような区間」はそれだけで簡単にSを超えてしまって解になりえない。具体的には8桁以上になると桁全体を覆う区間はSの上限108を超える( (108 - 107) * 8 ≒ 7×108 )。よってf(l)>=8のケースでは区間がある桁をまるまるまたぐということがありえないため、f(l)とf(r)の差はたかだか1にしかならない。これを元に以下のように場合分けをする。

(1) f(l) < 8 のとき

とりうる区間が結構面倒な形になりうるが、ここまでの大きさなら数列を陽に持った上で尺取りをやる形の全探索ができる。rの方は8桁の方に少し飛び出す可能性があることに注意する(厳密にやらずとも適当に5×107くらいまで持っておけばまあ足りる)。

(2) f(l) >= 8 のとき

区間の長さに注目すると、まずありうる区間長の範囲は1~S/8となる (f(l)の下限が8 => 長さ1の区間の最小値が8)。ここで区間長kを定めたとき、上述のとおりf(l)とf(r)の差はたかだか1であることから、f(l), f(l+1), ..., f(r) の並びは [S/k], [S/k], [S/k], [S/k]+1, [S/k]+1 のような形になる([]はここでは切り捨ての意)。つまり後半の(S%k)個が[S/k]+1, 前半の(K-S%K)個が[S/k]という形である。Sがkで割り切れないとき、このようなfの並びはただ一通りに定まる。よって1<=k<=S/8かつSの約数でないkの数だけ答えに+1が加わる。一方でSがkで割り切れるときはすべてが[S/k]となる。これはつまりf(l)=f(r)となるケースで、このときはその桁の中で好きに区間をとることができる。なのでこのケースをカバーするためにはSの約数を全部求めてf(l)を列挙すればよい。それぞれのf(l)について具体的な数は (l桁の数の個数 - 区間長 + 1) となる。このとき区間長が上述の範囲に収まっているかチェックする必要があることに注意する。

感想

説明ができなすぎる 解説放送を見て通したのでわかりたい人はそっちをみてくれ 小さいものだけ全探索で大きいときサボるといえば結構あるあるに聞こえるがサボり方がなんか自分にとってかなり非自明だった

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.string;

immutable long MOD = 10^^9 + 7;

long powmod(long a, long x, long m) {
    long ret = 1;
    while (x) {
        if (x % 2) ret = ret * a % m;
        a = a * a % m;
        x /= 2;
    }
    return ret;
}

long count(int x) {
    return ((powmod(10, x, MOD) - powmod(10, x-1, MOD)) % MOD + MOD) % MOD;
}

void main() {
    auto S = readln.chomp.to!int;
    long ans = 0;

    // f(l) <= 7
    {
        auto A = new int[](10^^7*5);
        for (int k = 1; k <= 8; ++k) {
            for (int i = 10^^(k-1); i < 10^^k && i < 10^^7*5; ++i) {
                A[i] = k;
            }
        }

        for (int l = 1, r = 1, cnt = 0; l < 10^^7; cnt -= A[l], l += 1) {
            if (r < l) r = l, cnt = 0;
            while (cnt < S) cnt += A[r], r += 1;
            if (cnt == S) (ans += 1) %= MOD;
        }
    }

    // f(l) >= 8
    {
        int[] factors;
        int T = S / 8;

        for (int i = 1; i * i <= S; ++i) {
            if (S % i != 0) continue;
            if (i * i != S) factors ~= S / i;
            factors ~= i;
        }

        foreach (f; factors) {
            if (S / f < 8) continue;
            (ans += count(S / f) - f) %= MOD;
            ans = (ans + MOD) % MOD;
        }

        (ans += T) %= MOD;
    }

    ans.writeln;
}

Codeforces Round #562 (Div. 1) : B. Good Triple

https://codeforces.com/contest/1168/problem/B

問題概要

各文字が0または1のどちらかからなる長さNの文字列Sが与えられる。1 <= l <= r <= N を満たす整数の組(l, r)のうち、Sの区間[l, r]の中に同じ文字が3文字以上等間隔で並んでいるものの数を求めよ。

N <= 3×105

解法

なんかこの条件を満たさないのはすごい難しいんじゃないかと考えて小さいNで実験コードを書いてみると10文字以上くらいになると必ず条件を満たす3文字が存在してしまうことがわかる。よって長さ10程度までの区間を全探索して条件を満たさないものを全選び方の総数から引けばよい。

感想

本番中、実験まではやっていたのになぜかそこから全探索にいかず(かなしいね)

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.string;

void main(){
    auto S = readln.chomp;
    auto N = S.length.to!int;
    long ans = 1L * N * (N + 1) / 2;

    foreach (l; 0..N) foreach (r; l..min(N, l+10)) {
        bool ok = true;
        foreach (i; l..r+1) for (int j = 1; i - j >= l && i + j <= r; ++j) {
            if (S[i-j] == S[i] && S[i] == S[i+j]) ok = false;
        }
        ans -= ok;
    }

    ans.writeln;
}

Educational Codeforces Round 64: D. 0-1-Tree

https://codeforces.com/problemset/problem/1156/D

問題概要

N頂点の木が与えられる。木の各辺は0または1のコストを持つ。単純パスで、かつ辺が「0個以上の0コスト辺」→「0個以上の1コスト辺」の順序で並んでいるパスをvalidと呼ぶ。頂点x->yへの単純パスがvalidとなるような頂点の組(x, y)の数を求めよ。

N <= 2×105

解法

UnionFindを2つ作って、「0コスト辺だけを見たときの連結成分」と「1コスト辺だけを見たときの連結成分」を別々に作る。このとき同じ連結成分に属している頂点同士のペアは明らかにvalidであるので、連結成分ごとにこれをカウントする。単純な掛け算だが、ペアの順番はどちらもとれることに注意する(<x,y>も<y,x>も取れる)。次に各頂点ごとに「その頂点が属す0コスト辺での連結成分」と「その頂点が属す1コスト辺での連結成分」を見ると、前者から後者へのパスは「今見ている頂点を境に使う辺が0から1へと切り替わる」ようなvalidパスとなり、これも集合サイズの単純な掛け算でカウントできる(今見ている頂点自身は数に含めないことに注意する)。このような数え方は明らかに漏れ・ダブリがない。よってこれらを合計したものが答え。

感想

なんか木DPを考え始めちゃってだめだった それでなくてもUFを2つ持つという発想ができたかどうか、、

参考にした記事

https://cinnamo-coder.hatenablog.com/entry/2019/05/02/021544

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.string;

void main() {
    auto N = readln.chomp.to!int;
    auto G = new Tuple!(int, int)[][](N);
    foreach (i; 0..N-1) {
        auto s = readln.split.map!(to!int);
        G[s[0]-1] ~= tuple(s[1]-1, s[2]);
        G[s[1]-1] ~= tuple(s[0]-1, s[2]);
    }
    
    auto uf0 = new UnionFind(N);
    auto uf1 = new UnionFind(N);
    foreach (i; 0..N) foreach (to; G[i]) (to[1] == 0 ? uf0 : uf1).unite(i, to[0]);

    long ans = 0;
    foreach (i; 0..N) if (uf0.find(i) == i) ans += 1L * -uf0.table[i] * (-uf0.table[i]-1);
    foreach (i; 0..N) if (uf1.find(i) == i) ans += 1L * -uf1.table[i] * (-uf1.table[i]-1);
    foreach (i; 0..N) ans += 1L * (-uf0.table[uf0.find(i)] - 1) * (-uf1.table[uf1.find(i)] - 1);

    ans.writeln;
}


class UnionFind {
    int N;
    int[] table;

    this(int n) {
        N = n;
        table = new int[](N);
        fill(table, -1);
    }

    int find(int x) {
        return table[x] < 0 ? x : (table[x] = find(table[x]));
    }

    void unite(int x, int y) {
        x = find(x);
        y = find(y);
        if (x == y) return;
        if (table[x] > table[y]) swap(x, y);
        table[x] += table[y];
        table[y] = x;
    }
}