Find Longest Subsequence with Element Difference at Most K in Python

Suppose we are given a list of numbers and another value k. This time our task is to find the length of the longest subsequence where the absolute difference between every adjacent element is at most k.

So, if the input is like nums = [5, 6, 2, 1, −6, 0, −1, k = 4, then the output will be 6.

To solve this, we will follow these steps −

• Define a function update() . This will take i, x

• i := i + n

• while i is non−zero, do

  • segtree[i] := maximum of segtree[i], x

  • i := i / 2

• Define a function query() . This will take i, j

• ans := −infinity

• i := i + n

• j := j + n + 1

• while i < j is non−zero, do

  • if i mod 2 is same as 1, then

    • ans := maximum of ans, segtree[i]

    • i := i + 1

  • if j mod 2 is same as 1, then

    • j := j − 1

    • ans := maximum of ans, segtree[j]

  • i := i / 2

  • j := j / 2

• return ans

• Now, in the main function, do the following −

• nums = [5, 6, 2, 1, −6, 0, −1]

• k = 4

• n = 2 to the power (logarithm base 2 of (length of (nums) + 1) + 1)

• segtree := [0] * 100000

• snums := sort the list nums

• index := make a collection where x: i for each index i and element x in (snums)

• ans := 0

• for each x in nums, do

  • lo := return the leftmost position from snums, where (x − k) can be inserted while maintaining the sorted order

  • hi := (the leftmost position from snums, where (x + k) can be inserted while maintaining the sorted order) − 1

  • count := query(lo, hi)

  • update(index[x], count + 1)

  • ans := maximum of ans, count + 1

• return ans

Let us see the following implementation to get better understanding −

Example

 Live Demo

import math, bisect
class Solution:
   def solve(self, nums, k):
      n = 2 ** int(math.log2(len(nums) + 1) + 1)
      segtree = [0] * 100000
      def update(i, x):
         i += n
         while i:
            segtree[i] = max(segtree[i], x)
            i //= 2
         def query(i, j):
            ans = −float("inf")
            i += n
            j += n + 1
            while i < j:
               if i % 2 == 1:
                  ans = max(ans, segtree[i])
                  i += 1
               if j % 2 == 1:
                  j −= 1
                  ans = max(ans, segtree[j])
               i //= 2
               j //= 2
            return ans
         snums = sorted(nums)
         index = {x: i for i, x in enumerate(snums)}
         ans = 0
         for x in nums:
            lo = bisect.bisect_left(snums, x − k)
            hi = bisect.bisect_right(snums, x + k) − 1
            count = query(lo, hi)
            update(index[x], count + 1)
            ans = max(ans, count + 1)
      return ans
ob = Solution()
print(ob.solve([5, 6, 2, 1, −6, 0, −1], 4))

Input

[5, 6, 2, 1, −6, 0, −1], 4

Output

6