shomanのブログ

ただの備忘録

AtCoder ABC165 C - Many Requirements

以下の条件を満たす数列Aを全て生成します。

  • Aは、長さNの正整数列である。
  • 1 ≤ A_1 ≤ A_2 ≤ \cdots ≤ A_N ≤ M

その後、それぞれに対しQ個ある制約を当てはめスコアを計算し、最大値を求めます。

atcoder.jp

上の条件を満たす数列の生成は深さ優先探索(DFS)でできます。

Goでの実装は以下です。

// 条件を満たす数列を全て生成
func dfs(seq []int, rst *[][]int) {
    if len(seq) == n {
        *rst = append(*rst, seq)
        return
    }

    // 1つ前の要素の数 (最初の要素の場合は1)
    var last int
    if len(seq) == 0 {
        last = 1
    } else {
        last = seq[len(seq)-1]
    }
    // last以上M以下の数を追加
    for i := last; i <= m; i++ {
        dfs(append(seq, i), rst)
    }
}

var n, m int

func main() {
    seq := make([]int, 0)
    rst := make([][]int, 0)
    dfs(seq, &rst)
    fmt.Println(rst)
    // n=m=3の場合: [[1 1 1] [1 1 2] [1 1 3] [1 2 2] [1 2 3] [1 3 3] [2 2 2] [2 2 3] [2 3 3] [3 3 3]]
}

あとはそれぞれの数列に対し、Q個の制約を当てはめ、スコアを計算していくだけです。

解答例の全体は以下のようになります。

type item struct {
    A, B, C, D int
}

var n, m int

// 条件を満たす数列を全て生成
func dfs(seq []int, rst *[][]int) {
    if len(seq) == n {
        *rst = append(*rst, seq)
        return
    }

    // 1つ前の要素の数 (最初の要素の場合は1)
    var last int
    if len(seq) == 0 {
        last = 1
    } else {
        last = seq[len(seq)-1]
    }
    // last以上M以下の数を追加
    for i := last; i <= m; i++ {
        dfs(append(seq, i), rst)
    }
}

// 数列の得点を制約より計算
func calcPoint(seq []int, items []item) int {
    point := 0
    for _, it := range items {
        if seq[it.B]-seq[it.A] == it.C {
            point += it.D
        }
    }
    return point
}

// 全ての数列の中から最大の得点を計算
func findMaxPoint(seqs [][]int, items []item) int {
    max := 0
    for _, seq := range seqs {
        if v := calcPoint(seq, items); v > max {
            max = v
        }
    }
    return max
}

func solve(items []item) int {
    seq := make([]int, 0)
    rst := make([][]int, 0)
    dfs(seq, &rst)
    return findMaxPoint(rst, items)
}

func main() {
    n, m = nextInt(), nextInt()
    q := nextInt()
    items := make([]item, q)
    for i := 0; i < q; i++ {
        items[i] = item{nextInt() - 1, nextInt() - 1, nextInt(), nextInt()}
    }
    ans := solve(items)
    fmt.Println(ans)
}

この問題は制約側を先に処理しようとするとハマると思います…