LC 1109. Corporate Flight Bookings

https://leetcode.com/problems/corporate-flight-bookings/


这题目一上来就是很直白的区间计数,我的第一个想法就是区间树。虽然实现起来很麻烦,但是想试试看看是否work. 因为如果实现正确的话,插入时间复杂度就是O(nlgn), 最后遍历时间复杂度是O(n). 对n<=20000的数据集合应该是没有问题的。

事实证明我实现的区间树是没有问题的,但是就是慢了一点3436ms.

class Tree:
    def __init__(self, s, e):
        self.s = s
        self.e = e
        self.left = None
        self.right = None
        self.val = 0


def build_tree(s, e):
    if s == e:
        return Tree(s, e)
    m = (s + e) // 2
    t1 = build_tree(s, m)
    t2 = build_tree(m + 1, e)
    t = Tree(s, e)
    t.left = t1
    t.right = t2
    return t


def update_tree(t, s, e, v):
    if t.s == s and t.e == e:
        t.val += v
        return

    m = (t.s + t.e) // 2
    # [t.s .. m]
    if t.s <= s <= m:
        update_tree(t.left, s, min(m, e), v)
    # [m+1.. t.e]
    if (m + 1) <= e <= t.e:
        update_tree(t.right, max(m + 1, s), e, v)

    return


def walk_tree(t, ans, pfx):
    if t.s == t.e:
        ans.append(pfx + t.val)
        return

    pfx += t.val
    walk_tree(t.left, ans, pfx)
    walk_tree(t.right, ans, pfx)
    return


class Solution:
    def corpFlightBookings(self, bookings: List[List[int]], n: int) -> List[int]:
        t = build_tree(1, n)
        for i, j, k in bookings:
            update_tree(t, i, j, k)
        ans = []
        walk_tree(t, ans, 0)
        return ans

按照这个时间来看,肯定是还有更好的算法。这个问题就是典型的进出区间的问题,我们可以在

因为时间点是同质的可累加,所以比如进入时间点t有n次的话,那么就可以在进入t的时候+n.

为了遍历方便,我们可以对时间点做排序。然后在遍历所有的时间点期间,检查对应的进出时间点是否有事件产生。

class Solution:
    def corpFlightBookings(self, bookings: List[List[int]], n: int) -> List[int]:
        eva = [(x[0], x[2]) for x in bookings]
        evb = [(x[1], x[2]) for x in bookings]

        eva.sort()
        evb.sort()

        ans = []
        res = 0
        pa, pb = 0, 0
        for i in range(1, n + 1):
            while pa < len(eva) and eva[pa][0] == i:
                res += eva[pa][1]
                pa += 1
            ans.append(res)

            while pb < len(evb) and evb[pb][0] == i:
                res -= evb[pb][1]
                pb += 1

            if pb == len(evb):
                ans.extend([res] * (n - i))
                break

        return ans

接着如果我们的视角从events转移到moments上的话,那么可以设计出更简单的算法。

class Solution:
    def corpFlightBookings(self, bookings: List[List[int]], n: int) -> List[int]:
        tmp = [0] * (n + 2)

        for i, j, k in bookings:
            tmp[i] += k
            tmp[j + 1] -= k

        ans = []
        res = 0
        for i in range(1, n + 1):
            res += tmp[i]
            ans.append(res)

        return ans