天下一プログラマーコンテスト 2019 D - Three Colors

• 問題
• AC

$S = \sum_i a_i$とせよ。Rの和が$S/2$以上だと三角形は作れないし、G、Bについても同様である。三角形を作れないような塗り方の総数は（Rの和が$S/2$以上になるような塗り方）+（Gの和が$S/2$以上になるような塗り方）+（Bの和が$S/2$以上になるような塗り方）でだいたい求まるので、それを$3^N$から引けば答えになる。「だいたい」というのは、$S$が偶数の時は色Xの和が$S/2$かつ色Yの和が$S/2$であるような塗り方が存在するかもしれないので、上の式だと重複して数えてしまうことになるからである。

とりあえず$S$が奇数の場合を考える。数列を$a_1, ..., a_x$に限定した時にRの和が$u$になるような塗り方の総数を$f(x, u)$とすると、

$f(x, u) = f(x-1, u-a_x) + 2f(x-1, u)$

である。G、Bについても同じ数になるので、答えは

$3^N - 3 \sum_{u = \lceil S/2 \rceil}^S f(N, u)$

であり、DPすれば良い。

$S$が偶数の場合、重複して数えた分を引かなければならない。数列$a_1, ..., a_x$をX、Yの２色に塗り分けるとして、色Xの和が$u$になるような塗り方の総数を$g(x, u)$とすると、

$g(x, u) = g(x-1, u-a_x) + g(x-1, u)$
が成り立つ。上の式では、Rの和が$S/2$かつGの和が$S/2$、Gの和が$S/2$かつBの和が$S/2$、Bの和が$S/2$かつRの和が$S/2$という３つのパターンを２回数えているので、重複の総数は$3g(N, S/2)$である。したがって、$S$が偶数の場合の答えは
$3^N - \Bigl( 3\sum_{u = S/2}^S f(N, u) - 3g(N, S/2)\Bigr)$
になる。計算量は${\mathcal O}(N^2 \max_i a_i)$。

考察の序盤は、まず数列$a_1, ..., a_x$の塗り分けでRの和が$u$、Gの和が$v$になるようなものの総数を$f(x, u, v)$として考えていた。しかし、これではもちろん間に合わないしデータ構造による高速化も思いつけず、なんとか$f(x, u)$だけで済まないだろうかと悩んだ末に上の方針に思いいたった。補集合に注目するのが習慣化していればもう少し早く正解にたどり着けたかもしれない。