E - This Message Will Self-Destruct in 5s

[hxrxchang]

https://atcoder.jp/contests/abc166/tasks/abc166_e

問題概要

1~Nまでの番号がついた人たちがいる。それぞれの身長がわかっている。 番号の差の絶対値と身長の和が等しいペアの個数を求めよ。

解き方

i, jの組み合わせを列挙すると O(N ** 2) になる。二分探索などで効率化できるか考えたがだめ。

求めるのは j - a = a[i] + a[j] となるi, jの組み合わせ数であることから、この式を変形して、 i + a[i] = j - a[j] とすることができる。 そうすると、i, j を独立に考えることができるので O(N) で考えることができる。 i + a[i], j - a[j] となる数の組み合わせを計算して、それをかけ合わせると i + a[i] = j - a[j] となる組み合わせ数を求めることができる。

func solve() {
    n := getInt()
    a := getInts()

    // 求めるのは j - a = a[i] + a[j] となるi, jの組み合わせ数
    // 変形すると i + a[i] = j - a[j] となる
    // i + a[i], j + a[j] の組み合わせ数を求める
    cntI := make(map[int]int)
    cntJ := make(map[int]int)

    for i := 0; i < n; i++ {
        cntI[i + a[i]]++
        cntJ[i - a[i]]++
    }

    //  i + a[i], j + a[j] の組み合わせ数をかけ合わせて、i + a[i] = j - a[j] となる組み合わせ数を求める
    ans := 0
    for i := 0; i < n; i++ {
        ans += cntI[i] * cntJ[i]
    }

    fmt.Println(ans)
}

https://atcoder.jp/contests/abc166/submissions/55498341