読者です 読者をやめる 読者になる 読者になる

euphonictechnologies’s diary

Haskell超初心者の日記です。OCamlが好きです。

follow us in feedly

TopCoder SRM 694 Div2 Hard : UpDownNess

プログラミングコンテスト

全然解けなかった。頑張って一週間かけて解いてみた。

問題

TopCoder Statistics - Problem Statement

1からNまで数字を並べるとき、並べ終わった列の中のlo-hi-loの出現数を考えます。lo-hi-loとは0<=i<j<kのときに並べ終わった列x中でx[i]<x[j]かつx[j]>x[k]となるような(要は真ん中がこんもりしてる)3つの数字の組です。例えば[ 1, 3, 4, 2 ] には(1, 3, 2), (1, 4, 2), (3, 4, 2)の3つのlo-hi-loがあります。

問題はNとKが与えられます。1からNまでを含む全ての列の中でlo-hi-loの数がK個になる列の数をカウントします。

方針

出題者のスライドがあるのでそちらを参照。DPで解く。

www.slideshare.net DPの配列は

dp[i][j] = 数列のパターン数

where

  • i = 1からiまでの数字を並べ終えたとき
  • j = j個のlo-hi-loを含む

当然答えはd[N][K]となります。なので、3重ループで

  1. iについて1からNまでループ、
  2. jについて0からKまでループ、
  3. i+1をiまでの数字の列に入れる時にどこの隙間に入れるか、例えば0なら先頭、3なら数列の3番めの要素の前、iなら一番最後です(ここ自分の少し失敗したところで、1からNまで並べるので、iが3の時、要素は3つで0オリジンの配列のインデックスと1オリジンの並べてる数字の最大数がずれます。なのでiの時一番最後となります)。

Code

3種類書いてみました。 まずは普通のforまわすDP。 lとrはそれぞれ挿入箇所の左にある要素の数、右にある要素の数です。

public class UpDownNess {
    public static final int DIVISOR = 1000000007;

    public int count(int N, int K) {
        int[][] dp = new int[5005][5005];
        dp[0][0] = 1;

        for (int i = 0; i < N + 1; i++) {
            for (int j = 0; j < K + 1; j++) {
                for (int l = 0; l <= i; l++) {
                    int r = i - l;
                    if (j >= l * r) {
                        dp[i + 1][j] += dp[i][j - l * r];
                        dp[i + 1][j] %= DIVISOR;
                    }
                }
            }
        }

        return dp[N][K] % DIVISOR;
    }
}

次はメモ化再帰。countInnerは右と左にある要素の数がわかるとき分割統治法的に求められる、という考えが強調されていて私としては理解しやすいです。4つめを入れる時には3つ目までの時のパターンがわかっていると良くて、それは左に何もないとき、左が1つ、右2つの時、左が2つ、右1つの時、右に何もない時の4パターンで、それぞれのパターン内で再帰的にこの考え方を適用します。

public class UpDownNess {
    private static final int DIVISOR = 1000000007;
    private static int[][] dp = new int[5005][5005];

    int countInner(int n, int k) {
        if (n == 1 && k == 0) return 1;
        if (dp[n][k] >= 0) return dp[n][k];

        int ret = 0;
        for (int l = 0; l < n; ++l) {
            int r = n - 1 - l;
            if (k >= l * r)
                ret += countInner(n - 1, k - l * r);
            ret %= DIVISOR;
        }
        System.out.printf("%d %d called\n", n, k);
        dp[n][k] = ret % DIVISOR;
        return dp[n][k];
    }

    public int count(int N, int K) {
        for (int i = 0; i < 5005; i++) {
            for (int j = 0; j < 5005; j++) {
                dp[i][j] = -1;
            }
        }
        return countInner(N, K) % DIVISOR;
    }
}

最後にJava8特有のFunctionを使った汎用memoizeでのメモ化再帰。メモ化する関数はInt -> Int -> Intなのでmemoizeを二回使ってcurryingみたいなことをしています。何も考えずHashMapを使いました。つまり配列です。TreeMapとか使うよりは早いかと。

import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

public class UpDownNess {
    public static final int DIVISOR = 1000000007;

    public static <I, O> Function<I, O> memoize(Function<I, O> f) {
        Map<I, O> lookup = new HashMap<>();
        return input -> lookup.computeIfAbsent(input, f);
    }

    static int countInner(int n, int k) {
        if (n == 1 && k == 0) return 1;

        int ret = 0;
        for (int l = 0; l < n; ++l) {
            int r = n - 1 - l;
            if (k >= l * r)
                ret += cached.apply(n - 1).apply(k - l * r);
            ret %= DIVISOR;
        }
        return ret % DIVISOR;
    }

    private static final Function<Integer, Function<Integer, Integer>> cached =
            memoize(x -> memoize(y -> countInner(x, y)));

    public int count(int N, int K) {
        return cached.apply(N).apply(K) % DIVISOR;
    }
}

countInner内でもcachedを呼び出さないと全然メモ化されなくて意味ないのでそこに気をつける(私は15分ほど無駄にしました)。

ちなみにmemoize関数を使った実装でもシステムテスト通ります。時間的にも問題なしなのでさくっとメモ化再帰したい時にはいいかも。memoize関数をコピって貼っつけるだけで使えるしいい感じです。感覚的には普通のdp配列を使ったものに対するオーバーヘッドは+10~20%ぐらい。

反省

なぜか完全な全列挙をやろうとしていてperm関数を実装したりcomb関数を実装したりしていた。流石に全列挙してフィルタしてカウントしていたら太陽が何度爆発してお姉さんが何人爆死しても足りないことぐらいは気づくべき。

あとperm関数がStreamの形でしかかけなくてほんとダメな方の関数脳だなって思った。気をつけたい(でも再帰とパターンマッチで2秒でかけるものをwhile使って書くとかどうかしムニョムニョ...)。

[追記: 20160717 01:44] 加筆しました。