Find kth element in an expanding string

Question

Given a string of the form AB2C3 and an int k.Expand the string as ABABC3 then ABABCABABCABABC. The task is to find the kth element. You have limited memory so you can not expand the entire string. You just have to find the kth element.

I am not sure about how to go about it. It was asked to my friend in a coding interview and I have given a lot of thought to it but I am not getting an efficient solution.

Answer 1

Update: An O(1) space and O(N) time version follows. See below.

---

Original solution uses O(1) space and O(N log k) time, where n is the size of the unexpanded string:

char find_kth_expanded(const char* s, unsigned long k) {
  /* n is the number of characters in the expanded string which we've
   * moved over.
   */
  unsigned long n = 0;
  const char *p = s;
  for (;;) {
    char ch = *p++;
    if (isdigit(ch)) {
      int reps = ch - '0';
      if (n * reps <= k)
        n *= reps;
      else {
        /* Restart the loop. See below. */
        k = k % n;
        p = s;
        n = 0;
      }
    }
    else if (ch == 0 || n++ == k)
      return ch;
  }
}

The function simply moves left to right through the string, keeping track of how many characters in the expanded string it has passed over. If that value reaches k, then we've found the kth character in the expanded string. If a repetition would skip over character k, then we reduce k to the index within the repetition, and restart the scan.

It's obvious that it uses O(1) space. To prove that it runs in O(N log k), we need to count the number of times that the loop is restarted. If we are restarting, then k≥n, because otherwise we would previously have returned the character at n. If k≥2n then n≤k/2 so k%n≤k/2. If k<2n then k%n = k-n. But n>k/2, so k-n<k-k/2 and therefore k%n<k/2.

So when we restart, the new value of k is at most half of the old value. So in the worst case, we'll restart log2k times.

---

While the above solution is easy to understand, we can actually do better. Instead of restarting the scan, we can just scan backwards once we scan past k (expanded) characters. During the backwards scan, we need to always correct k to the range in the current segment by taking its modulus base the segment length:

/* Unlike the above version, this one returns the point in the input
 * string corresponding to the kth expanded character.
 */
const char* find_kth_expanded(const char* s, unsigned long k) {
  unsigned long n = 0;
  while (*s && k >= n) {
    if (isdigit(*s))
      n *= *s - '0';
    else
      ++n;
    ++s;
  }
  while (k < n) {
    --s;
    if (isdigit(*s)) {
      n /= *s - '0';
      k %= n;
    }
    else
      --n;
  }
  return s;
}

Neither of the above functions correctly handle the case where the multiplier is 0 and k is less than the length of the segment multiplied by 0. If 0 is a legal multiplier, a simple solution would be to reverse scan the string for the last 0 and start find_kth_expanded at the following character. Since the reverse scan is O(N), the time complexity is not changed.