E - Distinct Adjacent

[hxrxchang]

https://atcoder.jp/contests/abc307/tasks/abc307_e

問題概要

1~Nの番号がついてN人の人が番号順に輪になって並んでいる。 N人の人に0以上M未満の整数を一つずつ渡す。 どの隣り合う二人が渡された数が異なるものの渡し方の数を求めよ。

解き方

1番に渡す色のパターンがM。2番に渡すパターンは1番に渡した色を除いてM-1、3番に渡すパターンは2番に渡した色を除いてM-1, ... としていくと、最後のN番がN-1番と1番に渡した色と同じかどうかを知らないと正しい通り数を導けない。 N番目のとき知りたいのは、"N-1番目に対して、1番目に渡した数字ではないものを渡したときのパターン数"。 こうすると、Nから1に繋げるときに、色が異なるので制約を満たすことができる。 N-1のときはN-2の... という感じで前の状態に依存していることが分かるのでDPで解く。

DPは、

• dp[i][0]: i番目まで配った際に、先頭の人と異なる数を渡したときのパターン数
• dp[i][1]: i番目まで配った際に、先頭の人と同じ数を渡したときのパターン数

を管理するとよい。

状態遷移の方法は、

• dp[i][0] (先頭の人と異なる数字を貰う) は、dp[i - 1][0] * (M - 2) + dp[i - 1][1] * (M - 1)
  • dp[i - 1][0] からの遷移で、M - 2 で-2するのは、先頭の人と違う数字かつ一個前の人と違う数字なのでMから2つ引く
  • dp[i - 1][1] からの遷移は、一個前の人とだけ違う数字ならいいので、その分1つだけ引く
• dp[i][1] (先頭の人と同じ数字を貰う) は、先頭の人と同じ数字を貰う一通りだけ。一個前の先頭とは違う数字をもらった場合から遷移する必要があるので、dp[i - 1][0]

func solve() {
    in := getInts()
    n, m := in[0], in[1]

    // dp[i][j]: i-1番目までの人に対して、
    // dp[i][0]: 先頭(i=0)の人と異なる数を渡したときのパターン数
    // dp[i][1]: 先頭(i=0)の人と同じ数を渡したときのパターン数
    dp := make([][]int, n)
    for i := 0; i < n; i++ {
        dp[i] = make([]int, 2)
    }
    // dp[0][0] は起こり得ないので0のまま
    dp[0][1] = m

    for i := 1; i < n; i++ {
        // 先頭に渡した数と異なるとき
        dp[i][0] = (dp[i - 1][1] * (m - 1) + dp[i - 1][0] * (m - 2)) % MOD

        // 先頭に渡した数と同じ時
        dp[i][1] = dp[i - 1][0]
    }

    fmt.Println(dp[n - 1][0])
}

https://atcoder.jp/contests/abc307/submissions/56386982