LintCode 840. 可变范围求和

2020-08-17  本文已影响0人  CW不要无聊的风格

题目描述

给定一个整数数组 nums, 然后你需要实现两个函数:

-- update(i, val) 将数组下标为i的元素修改为val

-- sumRange(l, r) 返回数组下标在[l, r][l,r]区间的元素的和

注意事项

数组只能通过update函数进行修改。

你可以假设 update 函数与 sumRange 函数的调用数量是均匀的。


样例

输入:

  nums = [1, 3, 5]

  sumRange(0, 2)

  update(1, 2)

  sumRange(0, 2)

输出:

  9

  8

输入:

  nums = [0, 9, 5, 7, 3]

  sumRange(4, 4)

  sumRange(2, 4)

  update(4, 5)

  update(1, 7)

  update(0, 8)

  sumRange(1, 2)

输出:

  3

  15

  12


思路

列表切片返回区间和是行不通的,别天真了,会超时,这题可以用线段树来解。

树节点代表的是一个区间[left, right],存储的是区间和,叶子节点存储的是数组中的每个数字。除叶子节点外,每个树节点都会划分2个子区间[left, mid], [mid+1, right]分配给其左、右儿子,其中mid就是left和right的中点,于是就有这个性质:每个节点的值是其左、右儿子的值之和。从而在更新节点值时,我们需要对应地将所有关联的父节点也更新。


代码

1. 

class NumArray:

    class TreeNode:

        def __init__(self, start, end, val):

            """每个节点存储的是由start~end位置的数字之和,

              左儿子是start~mid,右儿子是mid+1~end,

              叶子节点的start=end"""

            self.start = start

            self.end = end

            self.val = val

            self.left = None

            self.right = None

    def __init__(self, nums):

        """

        :type nums: List[int]

        """

        self.root = self._build_tree(nums, 0, len(nums) - 1)

    def update(self, i, val):

        """

        :type i: int

        :type val: int

        :rtype: void

        """

        self._modify(self.root, i, val)

    def sumRange(self, i, j):

        """

        :type i: int

        :type j: int

        :rtype: int

        """

        return self._query(self.root, i, j)

    def _build_tree(self, nums, start, end):

        node = self.TreeNode(start, end, nums[start])

        # 叶子节点

        if start == end:

            return node

        mid = (start + end) // 2

        node.left = self._build_tree(nums, start, mid)

        node.right = self._build_tree(nums, mid + 1, end)

        # 本节点存储的值等于其左儿子与右儿子存储的值之和

        node.val = node.left.val + node.right.val

        return node

    def _modify(self, node, i, val):

        """修改位置i的节点值"""

        # 已找到对应的叶子节点,直接更新节点值

        if node.start == node.end:

            node.val = val

            return

        if node.left.end < i:

            self._modify(node.right, i, val)

        else:

            self._modify(node.left, i, val)

        # 最后要更新本节点的值       

        node.val = node.left.val + node.right.val

    def _query(self, node, i, j):

        """查询位置区间[i,j]的数字之和"""

        # 若节点代表的区间被包含在查询区间内,则直接返回其存储的值

        if i <= node.start and node.end <= j:

            return node.val

        # 节点代表的区间与查询区间无交集

        if node.start > j or node.end < i:

            return 0

        # 节点代表的区间与查询区间有交集,则需要进一步在其左、右儿子中查询

        return self._query(node.left, i, j) + self._query(node.right, i, j)

2.

class NumArray:

    def __init__(self, nums):

        """

        :type nums: List[int]

        """

        self.size = len(nums)

        # nums中的数字用作叶子节点,同时对应构造等数量的父节点,

        # 用于存储区间和

        self.tree = [0] * self.size + nums

        for i in range(self.size - 1, 0, -1):

            # 左儿子2i、右儿子2i+1

            self.tree[i] = self.tree[i << 1] + self.tree[i << 1 | 1]

    def update(self, i, val):

        """

        :type i: int

        :type val: int

        :rtype: void

        """

        i += self.size

        # 先更新对应的叶子节点

        self.tree[i] = val

        while i > 1:

            # 当i是偶数时,代表左儿子,i XOR 1则为i+1,即右儿子;

            # 当i是奇数时,代表右儿子,i XOR 1 则为i-1,即左儿子

            self.tree[i >> 1] = self.tree[i] + self.tree[i ^ 1]

            i >>= 1

    def sumRange(self, i, j):

        """

        :type i: int

        :type j: int

        :rtype: int

        """

        # 从叶子节点回溯

        i += self.size

        j += self.size

        s = 0

        while i <= j:

            # i是奇数,代表右儿子

            if i & 1:

                s += self.tree[i]

                # 变为后一个子区间的左儿子

                i += 1

            # 后一个子区间的父节点

            i >>= 1

            # j是偶数,代表左儿子

            if not j & 1:

                s += self.tree[j]

                # 变为前一个子区间的右儿子

                j -= 1

            # 前一个子区间的父节点

            j >>= 1

        return s

上一篇下一篇

猜你喜欢

热点阅读