AtCoder Regular Contest 125 D問題メモ

D - Unique Subsequence

問題

長さ $N$ の数列 $A=(A_1,A_2,...,A_N)$
$A$ の（連続するとは限らない）部分列のうち、以下の条件を満たすものの個数を求めよ
条件
- 各要素を $A$ のどの位置から取り出したかが一意に特定できる
$1 \le N \le 2 \times 10^5$

例

A = (3 1 4 1 5)

部分列 (1 4)  : OK
       (1 5)  : NG 1を取り出した位置が2文字目か4文字目か特定できない
       (1 1 5): OK

解法

問題内容と制約からDPっぽいのは察せるが、どこから遷移できて、どこから遷移してはいけないかを明確にしないと混乱する。

$DP[i][j]=i$ 番目までを考えて、$A_j$ を最後に使ったもののうち、条件を満たしているパターン数

実際は $i$ の次元は省略して $j$ だけで実装する。

また、indexは0からに読み替えて説明する。

基本的な更新

$i$ の昇順に更新していく。

$i$ を進めても、$j$ は基本的に $DP[i][j] = DP[i-1][j]$ となり引き継がれるので、$i$ の次元を省略していたら特に何もしなくていい。

$DP[i][i]$ だけ新しく計算する。

$\displaystyle DP[i][i] = \sum_{j=0}^{i-1}DP[i-1][j]$ + 1

最後に使った位置が $j=0～i-1$ からそれぞれ遷移できて、また $A_i$ を先頭とする部分列が新たに1つ加わる。

これを計算していくと $1,2,4,8,...$ と2の冪乗となる。仮に $A$ に出てくる要素に重複が無ければ、この総和が答えとなる。

どこから遷移してはいけないか

重複があると、一意に特定できない場合が出てくる。

i  0 1 2 3 4 5 6 7 8 9
A  3 1 4 1 5 9 2 6 5 3
                   ^ i=8 において、この5を使う場合、
   ~~~~~~~ ^ i=4 より手前から遷移すると、この5と区別できなくなる
             また、A8を先頭としても同様に区別できなくなる

そのため、過去に同じ値が出てきた要素については、直前に同じ値が出現したindexを $p$ として、$p$ 未満から遷移できない。

$\displaystyle DP[i][i] = \sum_{j=p}^{i-1}DP[i-1][j]$

$j$ の和を取る開始indexが変わり、また最後の +1 もなくなる。

どこへ遷移してはいけないか

上記だけだと、$(5,3)$ のような条件を満たさないケースが数えられてしまう。

つまり、$i=9$ において $DP[9][9]$ を求めるとき、今のままだと波線部全体の $DP[8][j]$ の総和を取ることになる。

i  0 1 2 3 4 5 6 7 8 9
A  3 1 4 1 5 9 2 6 5 3
   ~~~~~~~~~~~~~~~~~

しかしこの中には複数の同じ値 $5$ などが挟まれていて、このままだとどちらの $5$ が最後だったケースからも遷移されてしまい、区別できなくなる。

同じ値に関しては最後に出現した位置、この場合は $5$ に関して $j=4$ は除外して $j=8$ からのみ遷移したい。

i  0 1 2 3 4 5 6 7 8 9
A  3 1 4 1 5 9 2 6 5 3
   ~   ~~~   ~~~~~~~

要は、$DP[○][j]$ の値は、$A_j$ と同じ値が次に出てくる位置を $n$ として、$DP[n+1以降][n+1以降]$ に遷移してはいけない。

i  0 1 2 3 4 5 6 7 8 9 10 11
A  3 1 4 1 5 9 2 6 5 3  5  8
           ^ i=4のDPの値は
                   ^ xxxxxxx 次に同じ値が出てくるi=8より後に遷移してはいけない

$i=8$ の時点で、$DP[8][8]$ の値を計算した後に $DP[8][4]$ の値を0に更新しておけば、これを達成できる。

$DP[i][p]=0$ （※$A_i$ と同じ値が以前に出現している場合）

まとめ

$A_i$ がはじめて出現する値の場合
- $\displaystyle DP[i][i] = \sum_{j=0}^{i-1}DP[i-1][j]$ + 1
$A_i$ が2回目以降に出現知る値の場合、直前に出現したindexを $p$ として
- $\displaystyle DP[i][i] = \sum_{j=p}^{i-1}DP[i-1][j]$
- $DP[i][p]=0$

これで $DP$ が正しく求められ、最後まで処理したら $DP[N-1]$ の総和が答えとなる。

区間の和を頻繁に求めるため、DP配列自体を Segment Tree や Fenwick Tree などで実装すれば、$O(N \log{N})$ となる。

配るDPともらうDPの両面から条件に合わないものの除外を行う感じで、一方だけで考えているとなかなかたどり着かなかった。

Python3

class SegmentTreeInjectable:
    def __init__(self, n, identity_factory, func):
        n2 = 1 << (n - 1).bit_length()
        self.offset = n2
        self.tree = [identity_factory() for _ in range(n2 << 1)]
        self.func = func
        self.idf = identity_factory

    @classmethod
    def from_array(cls, arr, identity_factory, func):
        """ 既存の配列から生成 """
        ins = cls(len(arr), identity_factory, func)
        ins.tree[ins.offset:ins.offset + len(arr)] = arr
        for i in range(ins.offset - 1, 0, -1):
            l = i << 1
            r = l + 1
            ins.tree[i] = func(ins.tree[l], ins.tree[r])
        return ins

    def add(self, i, x):
        """
        Aiにxを加算
        :param i: index (0-indexed)
        :param x: add value
        """
        i += self.offset
        self.tree[i] = self.func(self.tree[i], x)
        self.__upstream(i)

    def update(self, i, x):
        """
        Aiの値をxに更新
        :param i: index(0-indexed)
        :param x: update value
        """
        i += self.offset
        self.tree[i] = x
        self.__upstream(i)

    def __upstream(self, i):
        tree = self.tree
        func = self.func
        while i > 1:
            tree[i >> 1] = func(tree[i], tree[i ^ 1])
            i >>= 1

    def get_range(self, a, b):
        """
        [a, b)の値を得る
        :param a: index(0-indexed)
        :param b: index(0-indexed)
        """
        tree = self.tree
        func = self.func
        result_l = self.idf()
        result_r = self.idf()

        l = a + self.offset
        r = b + self.offset
        while l < r:
            if r & 1:
                result_r = func(tree[r - 1], result_r)
            if l & 1:
                result_l = func(result_l, tree[l])
                l += 1
            l >>= 1
            r >>= 1

        return func(result_l, result_r)

    def get_all(self):
        return self.tree[1]

    def get_point(self, i):
        return self.tree[i + self.offset]

    def debug_print(self):
        i = 1
        while i <= self.offset:
            print(self.tree[i:i * 2])
            i <<= 1


n = int(input())
aaa = list(map(int, input().split()))
MOD = 998244353

prev_index = [-1] * n
prev_index_tmp = {}
for i, a in enumerate(aaa):
    if a in prev_index_tmp:
        prev_index[i] = prev_index_tmp[a]
    prev_index_tmp[a] = i

sgt = SegmentTreeInjectable(n, int, lambda a, b: (a + b) % MOD)

for i, a in enumerate(aaa):
    p = prev_index[i]

    if p == -1:
        sgt.update(i, sgt.get_range(0, i) + 1)
    else:
        sgt.update(i, sgt.get_range(p, i))
        sgt.update(p, 0)

ans = sgt.get_all()
print(ans)

目次

AtCoder Regular Contest 125 D問題メモ

D - Unique Subsequence

問題

例

解法

基本的な更新

どこから遷移してはいけないか

どこへ遷移してはいけないか

まとめ