LeetCode-Solutions/Python/sum-of-beautiful-subsequences.py at master · kamyu104/LeetCode-Solutions · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
# Time:  precompute: O(rlogr), r = max_nums
#        runtime:    O(mx + nlogr * (log(nlogr) + logn)), mx = max(nums)
# Space: O(rlogr)

# number theory, bit, fenwick tree
MOD = 10**9+7
class BIT(object):  # 0-indexed.
    def __init__(self, n):
        self.__bit = [0]*(n+1)  # Extra one for dummy node.

    def add(self, i, val):
        i += 1  # Extra one for dummy node.
        while i < len(self.__bit):
            self.__bit[i] = (self.__bit[i]+val) % MOD
            i += (i & -i)

    def query(self, i):
        i += 1  # Extra one for dummy node.
        ret = 0
        while i > 0:
            ret = (ret+self.__bit[i]) % MOD
            i -= (i & -i)
        return ret


def factors(n):  # Time: O(nlogn)
    result = [[] for _ in xrange(n+1)]
    for i in xrange(1, n+1):
        for j in range(i, n+1, i):
            result[j].append(i)
    return result


def phi_sieve(n):  # Time: O(nlog(logn))
    phi = range(n+1)
    for i in xrange(2, n+1):
        if phi[i] != i:
            continue
        for j in xrange(i, n+1, i):
            phi[j] -= phi[j]//i
    return phi


MAX_NUM = 7 * 10**4
FACTORS = factors(MAX_NUM)
PHI = phi_sieve(MAX_NUM)
class Solution(object):
    def totalBeauty(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        def count(arr):
            for i, x in enumerate(sorted(arr)):  # coordinate compression
                val_to_idx[x] = i
            bit = BIT(len(arr))
            for x in arr:
                bit.add(val_to_idx[x], bit.query(val_to_idx[x]-1)+1)
            return bit.query(len(arr)-1)

        mx = max(nums)
        val_to_idx = [0]*(mx+1)
        lookup = [[] for _ in xrange(mx+1)]
        for x in nums:
            for d in FACTORS[x]:
                lookup[d].append(x)
        return reduce(lambda accu, x: (accu+x)%MOD, (PHI[g]*count(lookup[g]) for g in reversed(xrange(1, mx+1))), 0)


# Time:  precompute: O(rlogr), r = max_nums
#        runtime:    O(mx * log(mx) + nlogr * (log(nlogr) + logn)), mx = max(nums)
# Space: O(rlogr)
# number theory, bit, fenwick tree
class Solution2(object):
    def totalBeauty(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        def count(arr):
            val_to_idx = {x:i for i, x in enumerate(sorted(set(arr)))}  # coordinate compression
            bit = BIT(len(val_to_idx))
            for x in arr:
                bit.add(val_to_idx[x], bit.query(val_to_idx[x]-1)+1)
            return bit.query(len(val_to_idx)-1)

        mx = max(nums)
        lookup = [[] for _ in xrange(mx+1)]
        for x in nums:
            for d in FACTORS[x]:
                lookup[d].append(x)
        result = 0
        cnt = [0]*(mx+1)
        for g in reversed(xrange(1, mx+1)):
            cnt[g] = count(lookup[g])
            for ng in xrange(g+g, mx+1, g):
                cnt[g] -= cnt[ng]
            result = (result+g*cnt[g])%MOD
        return result