あるエンジニアのAtCoder奮闘記

東京都港区にあるアミフィアブル株式会社のエンジニアが、AtCoderで解いた問題について振り返ったりしていく会社公認のブログです。

AtCoder Tenka1 Programmer Contest C - Align (Go)

irisruneです。お盆休みのある人もいたり学生さんは夏休みだったりしますが、よい休暇を過ごせるといいですね。

問題

atcoder.jp

なんとなくでも解法が見えそうな問題ですが、証明が難しかったり落とし穴があったりします。

package main

import (
    "bufio"
    "fmt"
    "os"
    "sort"
    "strconv"
)

var sc *bufio.Scanner

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

func intAbs(a int) int {
    if a > 0 {
        return a
    }
    return -a
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n := nextInt()

    aArray := make([]int, n)
    for i := range aArray {
        aArray[i] = nextInt()
    }

    ans1 := 0
    sort.Sort(sort.IntSlice(aArray))
    aSlice := aArray[:]
    aLeft, aRight := aSlice[0], aSlice[0]
    aSlice = aSlice[1:]
    for {
        if len(aSlice) == 0 {
            break
        }
        ans1 += intAbs(aLeft - aSlice[len(aSlice)-1])
        aLeft = aSlice[len(aSlice)-1]
        aSlice = aSlice[:len(aSlice)-1]
        if len(aSlice) == 0 {
            break
        }
        ans1 += intAbs(aRight - aSlice[len(aSlice)-1])
        aRight = aSlice[len(aSlice)-1]
        aSlice = aSlice[:len(aSlice)-1]

        if len(aSlice) == 0 {
            break
        }
        ans1 += intAbs(aLeft - aSlice[0])
        aLeft = aSlice[0]
        aSlice = aSlice[1:]
        if len(aSlice) == 0 {
            break
        }
        ans1 += intAbs(aRight - aSlice[0])
        aRight = aSlice[0]
        aSlice = aSlice[1:]
    }

    ans2 := 0
    sort.Sort(sort.Reverse(sort.IntSlice(aArray)))
    aSlice = aArray[:]
    aLeft, aRight = aSlice[0], aSlice[0]
    aSlice = aSlice[1:]
    for {
        if len(aSlice) == 0 {
            break
        }
        ans2 += intAbs(aLeft - aSlice[len(aSlice)-1])
        aLeft = aSlice[len(aSlice)-1]
        aSlice = aSlice[:len(aSlice)-1]
        if len(aSlice) == 0 {
            break
        }
        ans2 += intAbs(aRight - aSlice[len(aSlice)-1])
        aRight = aSlice[len(aSlice)-1]
        aSlice = aSlice[:len(aSlice)-1]

        if len(aSlice) == 0 {
            break
        }
        ans2 += intAbs(aLeft - aSlice[0])
        aLeft = aSlice[0]
        aSlice = aSlice[1:]
        if len(aSlice) == 0 {
            break
        }
        ans2 += intAbs(aRight - aSlice[0])
        aRight = aSlice[0]
        aSlice = aSlice[1:]
    }

    switch {
    case ans1 > ans2:
        fmt.Println(ans1)
    default:
        fmt.Println(ans2)
    }
}

直感的に考えると、一列に並べた時に中央に一番小さな要素が、その両隣に二番目までに大きな要素が、その両隣に三番目までに小さな要素が…という形で並べるとよさそうに思えます。与えられた入力例で考えると、

2 8 1 6 3

4 1 9 1 5 3

5 1 5

という並び順になります。なお、入力例2のようにNが偶数の場合は最後に残った要素を左右どちらにつけるかで結果が変わってしまうので、より大きな(小さな)値を左側か右側どちらかに優先的に付けるようにするとよいです。

この考え方は厳密に考えるとそれほど間違ってはおらず、公式解説の考え方を借りれば左右端の要素が与えられた整数群の中央値付近であればよかったりします。

ただし一つだけ落とし穴があり、例えば入力例3を少し変えたケースで

5 1 1

と出力してしまうと明らかに誤った結果となります。要するに、一番小さな要素を中央に置くケースと一番大きな要素を中央に置く(その両隣に小さい要素、大きい要素、…と置く)ケースとで結果が異なるわけです。実際には入力例2では結果が同じとなるため、Nが偶数の場合は気にしなくていいのかもしれません。

以上を踏まえて、2通りの整数の並べ方を試して結果が大きい方を出力するという形で解くことができます。

余談ですが、上記コードではGoのスライスの性質を活用しています。スライス上で範囲の切り出しを行っても元の配列は変わらないので、同じ配列の実体を参照しながら範囲の切り出しと復元を(実現しているかのように)実装できるわけですね。aArrayという名前で宣言しているのが既にスライスなのはご愛敬です。

雑記

  • 実はコンテスト参加当時解けなかった問題でした。
    • 代わりにD問題が解けてしまったので試合に勝って勝負に負けたというフレーズがぴったりでした。