LeetCode 2389: Longest Subsequence With Limited Sum

link

If we had nums sorted, we could compute its list of cumulative sums. For a query, its insert position in this list of cumulative sums would give the max length of subsequence for that query.

Time: \mathcal{O}( \max{m, n} \cdot \lg{n} ). Space: \mathcal{O}(m).

class Solution:
    def find_insert_position(self, x, nums) -> int:
        lo, hi = 0, len(nums)-1
        while lo <= hi:
            mid = (lo+hi) // 2
            if x == nums[mid]:
                return mid
            if x < nums[mid]:
                hi = mid-1
            else:
                lo = mid+1
        
        return lo

    def answerQueries(self, nums: List[int], queries: List[int]) -> List[int]:
        nums.sort()
        m, n = len(queries), len(nums)
        cumsum = [nums[0]] + [0]*(n-1)
        for i, x in enumerate(nums[1:], start=1):
            cumsum[i] = cumsum[i-1] + x

        ans = [0]*m
        for i, q in enumerate(queries):
            pos = self.find_insert_position(q, cumsum)
            if pos == len(cumsum) or q != cumsum[pos]:
                ans[i] = pos
            else:
                ans[i] = pos+1
        return ans

Leave a comment