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

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

AtCoder ABC 128 C - Switches (Go)

irisruneです。今日の記事は参加しなかった方のコンテストです。

atcoder.jp

ABC127とは打って変わって少々厄介な問題です。

package main

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

var sc *bufio.Scanner

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

type light struct {
    ss int
    p  int
}

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

    lights := make([]light, m)
    for i := range lights {
        k := nextInt()
        for j := 0; j < k; j++ {
            s := nextInt()
            lights[i].ss += 1 << uint(s-1)
        }
    }
    for i := range lights {
        lights[i].p = nextInt()
    }

    ans := 0
    for x := 0; x < int(math.Pow(2.0, float64(n))); x++ {
        flg := true
        for _, l := range lights {
            elec := l.ss & x
            cnt := 0
            for i := 0; i < n; i++ {
                cnt += (elec >> uint(i)) & 1
            }
            if cnt%2 != l.p {
                flg = false
            }
        }
        if flg {
            ans++
        }
    }
    fmt.Println(ans)
}

制約条件が緩いので単純な全探索問題ではあります。今回は練習も兼ねてbit全探索で実装しましたがDFSでもいいですね。

ただし電球が点灯する条件が少々複雑で、簡潔に書く方法がなかなか思いつきませんでした。今回は、

  1. 予め各電球に対して対応するスイッチの組み合わせを2進数変数に変換しておく。
  2. onになっているスイッチの組み合わせと1.で用意した電球に対応するスイッチの組み合わせとでAND演算した結果を取る。
  3. 2.の結果について1となるbit数を数え上げ、2で割った余りを求めて電球が点灯するか判定する。
  4. 3.をすべての電球に対して判定し、すべて点灯するなら解に1を足す。
  5. onになっているスイッチの組み合わせすべてに対して2.から4.を繰り返す。

といった実装の方法をとりましたが、もっと簡潔に書ける気はしています。

全探索が思いついて実装できるかと、細かい処理を正確に実装できるかの両方が問われるのでなかなか難しい問題だと思います。前回のC問題がかなり易しかっただけとも言えそうです。

雑記

  • 今回一番欲しかったのは変数名のセンスです。