Ashton String

gistfile1.txt

typedef string::size_type ssize;
typedef pair<ssize, ssize> valLenPair;

ssize get(uint64_t x) {
    double v = (sqrt((8. * x) + 1.) - 1)/2;
    double c = ceil(v);
    double dummy;
    double frac = modf(v, &dummy);
    double res = floor((c+1) * frac);
    return (ssize)res;
}

class RedBlackTreeSet {
public:
    RedBlackTreeSet(string& s): s(s) {
        len = s.length();
        root = NIL = new Node();
        NIL->parent = NIL->left = NIL->right = NIL;
        for (ssize i = 0; i < len; ++i) {
            insert(i);
        }
    }

    ~RedBlackTreeSet() {
        destroyTree(root);
        delete NIL;
    }

    void insert(const ssize& data) {
        Node *y = NIL, *x = root;
        pair<int, ssize> cmpRes = make_pair(0, 0);
        ssize upperLimitCmpLen = 0, lowerLimitCmpLen = 0;
        Node *upperLimit = NIL;
        bool upperLimitSet = false, lowerLimitSet = false;
        while (x != NIL) {
            y = x;
            cmpRes = compare(x->data.first, data);
            if (cmpRes.first > 0) {
                upperLimitCmpLen = cmpRes.second;
                upperLimit = x;
                upperLimitSet = true;
            } else if (cmpRes.first < 0) {
                lowerLimitCmpLen = cmpRes.second;
                lowerLimitSet = true;
            } else {
                return;
            }
            x = x->getSubNode(cmpRes.first > 0);
        }
        Node* z = new Node(data);
        z->left = z->right = NIL;
        if (upperLimitSet) {
            upperLimit->data.second = upperLimitCmpLen;
        }
        if (lowerLimitSet) {
            z->data.second = lowerLimitCmpLen;
        }
        z->parent = y;
        if (y == NIL) {
            root = z;
        } else {
            cmpRes = compare(y->data.first, z->data.first);
            y->getSubNode(cmpRes.first > 0) = z;
        }
        insertFixup(z);
    }
    
private:
    enum class Color { RED, BLACK };

    struct Node {
        valLenPair data;
        Color color;
        Node* parent;
        Node* left;
        Node* right;
        Node*& getSubNode(bool isLeft) { return isLeft ? left : right; }
        Node() : data(0, 0), color(Color::BLACK), parent(nullptr), left(nullptr), right(nullptr) {}
        Node(const ssize& data) : data(data, 0), color(Color::RED), parent(nullptr), left(nullptr), right(nullptr) {}
    };
    
    Node* treeMinimum(Node* node) {
        while (node->left != NIL) {
            node = node->left;
        }
        return node;
    }

public:
    class iterator {
    private:
        Node* current;
        Node* NIL;
        RedBlackTreeSet* tree;

    public:
        iterator(Node* node, Node* NIL, RedBlackTreeSet* tree) : current(node), NIL(NIL), tree(tree) {}

        iterator& operator++() {
            if (current->right != NIL) {
                current = tree->treeMinimum(current->right);
            } else {
                Node* y = current->parent;
                while (y != NIL && current == y->right) {
                    current = y;
                    y = y->parent;
                }
                current = y;
            }
            return *this;
        }

        pair<ssize, ssize>& operator*() {
            return current->data;
        }

        bool operator==(const iterator& other) const {
            return current == other.current;
        }

        bool operator!=(const iterator& other) const {
            return !(*this == other);
        }
    };
    
    iterator begin() {
        return iterator(treeMinimum(root), NIL, this);
    }

    iterator end() {
        return iterator(NIL, NIL, this);
    }

private:
    string& s;
    ssize len;

    Node* root;
    Node* NIL;

    void rotate(Node* x, bool isLeft = true) {
        Node* y = x->getSubNode(!isLeft);
        x->getSubNode(!isLeft) = y->getSubNode(isLeft);
        if (y->getSubNode(isLeft) != NIL) {
            y->getSubNode(isLeft)->parent = x;
        }
        y->parent = x->parent;
        if (x->parent == NIL) {
            root = y;
        } else {
            x->parent->getSubNode(x == x->parent->getSubNode(true)) = y;
        }
        y->getSubNode(isLeft) = x;
        x->parent = y;
    }

    void insertFixup(Node* z) {
        while (z->parent->color == Color::RED) {
            bool isLeft = (z->parent == z->parent->parent->left);
            Node* y = z->parent->parent->getSubNode(!isLeft);
            if (y->color == Color::RED) {
                z->parent->color = Color::BLACK;
                y->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->getSubNode(!isLeft)) {
                    z = z->parent;
                    rotate(z, isLeft);
                }
                z->parent->color = Color::BLACK;
                z->parent->parent->color = Color::RED;
                rotate(z->parent->parent, !isLeft);
            }
        }
        root->color = Color::BLACK;
    }

    pair<int, ssize> compare(ssize a, ssize b) {
        ssize cmp_len = s.length() - max(a, b);
        auto lhs = s.data() + a;
        auto rhs = s.data() + b;
        ssize count = 0;
        while (count < cmp_len && *lhs == *rhs) {
          ++count;
          ++lhs;
          ++rhs;
        }
        return make_pair((count < cmp_len) ? (*lhs - *rhs) : 0, count);
    };

    void destroyTree(Node* node) {
        if (node != NIL) {
            destroyTree(node->left);
            destroyTree(node->right);
            delete node;
        }
    }
};

char ashtonString(string s, uint64_t k) {
    RedBlackTreeSet prefixes(s);
    uint64_t p = 0;
    uint64_t lastP = 0;
    uint64_t adjustedK = k - 1;
    for (RedBlackTreeSet::iterator it = prefixes.begin(); it != prefixes.end(); ++it) {
        uint64_t lenShortest = (*it).second;
        uint64_t lenLongest = s.length() - (*it).first;
        uint64_t Sn = (((lenLongest + lenShortest + 1L) % 2L) == 0) ? (((lenLongest + lenShortest + 1L) / 2L) * (lenLongest - lenShortest)) : ((lenLongest + lenShortest + 1L) * (((lenLongest - lenShortest) + 1L) / 2L));
        p += Sn;
        if (adjustedK < p) {
            uint64_t adjustment = (((*it).second % 2L) == 0) ? ((*it).second / 2L) * ((*it).second + 1L) :  (*it).second * (((*it).second + 1L) / 2L);
            return s[(*it).first + get(adjustedK - lastP + adjustment)];
        }
        lastP = p;
    }
    return 0;
}