Task: Color quantization

C: typedef struct oct_node_t oct_node_t, *oct_node;
struct oct_node_t{
	
	uint64_t r, g, b;
	int count, heap_idx;
	oct_node kids[8], parent;
	unsigned char n_kids, kid_idx, flags, depth;
};


inline int cmp_node(oct_node a, oct_node b)
{
	if (a->n_kids < b->n_kids) return -1;
	if (a->n_kids > b->n_kids) return 1;

	int ac = a->count * (1 + a->kid_idx) >> a->depth;
	int bc = b->count * (1 + b->kid_idx) >> b->depth;
	return ac < bc ? -1 : ac > bc;
}


oct_node node_insert(oct_node root, unsigned char *pix)
{
#	define OCT_DEPTH 8
	

	unsigned char i, bit, depth = 0;
	for (bit = 1 << 7; ++depth < OCT_DEPTH; bit >>= 1) {
		i = !!(pix[1] & bit) * 4 + !!(pix[0] & bit) * 2 + !!(pix[2] & bit);
		if (!root->kids[i])
			root->kids[i] = node_new(i, depth, root);

		root = root->kids[i];
	}

	root->r += pix[0];
	root->g += pix[1];
	root->b += pix[2];
	root->count++;
	return root;
}


oct_node node_fold(oct_node p)
{
	if (p->n_kids) abort();
	oct_node q = p->parent;
	q->count += p->count;

	q->r += p->r;
	q->g += p->g;
	q->b += p->b;
	q->n_kids --;
	q->kids[p->kid_idx] = 0;
	return q;
}


void color_replace(oct_node root, unsigned char *pix)
{
	unsigned char i, bit;

	for (bit = 1 << 7; bit; bit >>= 1) {
		i = !!(pix[1] & bit) * 4 + !!(pix[0] & bit) * 2 + !!(pix[2] & bit);
		if (!root->kids[i]) break;
		root = root->kids[i];
	}

	pix[0] = root->r;
	pix[1] = root->g;
	pix[2] = root->b;
}


void color_quant(image im, int n_colors)
{
	int i;
	unsigned char *pix = im->pix;
	node_heap heap = { 0, 0, 0 };

	oct_node root = node_new(0, 0, 0), got;
	for (i = 0; i < im->w * im->h; i++, pix += 3)
		heap_add(&heap, node_insert(root, pix));

	while (heap.n > n_colors + 1)
		heap_add(&heap, node_fold(pop_heap(&heap)));

	double c;
	for (i = 1; i < heap.n; i++) {
		got = heap.buf[i];
		c = got->count;
		got->r = got->r / c + .5;
		got->g = got->g / c + .5;
		got->b = got->b / c + .5;
		printf("%2d | %3llu %3llu %3llu (%d pixels)\n",
			i, got->r, got->g, got->b, got->count);
	}

	for (i = 0, pix = im->pix; i < im->w * im->h; i++, pix += 3)
		color_replace(root, pix);

	node_free();
	free(heap.buf);
}




Go: package main

import (
    "container/heap"
    "image"
    "image/color"
    "image/png"
    "log"
    "math"
    "os"
    "sort"
)

func main() {
    f, err := os.Open("Quantum_frog.png")
    if err != nil {
        log.Fatal(err)
    }
    img, err := png.Decode(f)
    if ec := f.Close(); err != nil {
        log.Fatal(err)
    } else if ec != nil {
        log.Fatal(ec)
    }
    fq, err := os.Create("frog16.png")
    if err != nil {
        log.Fatal(err)
    }
    if err = png.Encode(fq, quant(img, 16)); err != nil {
        log.Fatal(err)
    }
}


func quant(img image.Image, nq int) image.Image {
    qz := newQuantizer(img, nq) 
    qz.cluster()                
    return qz.Paletted()        
}


type quantizer struct {
    img image.Image 
    cs  []cluster   
    px  []point     
    ch  chValues    
    eq  []point     
}

type cluster struct {
    px       []point 
    widestCh int     
    chRange  uint32  
}

type point struct{ x, y int }
type chValues []uint32
type queue []*cluster

const (
    rx = iota
    gx
    bx
)

func newQuantizer(img image.Image, nq int) *quantizer {
    b := img.Bounds()
    npx := (b.Max.X - b.Min.X) * (b.Max.Y - b.Min.Y)
    
    qz := &quantizer{
        img: img,
        ch:  make(chValues, npx),
        cs:  make([]cluster, nq),
    }
    
    c := &qz.cs[0]
    px := make([]point, npx)
    c.px = px
    i := 0
    for y := b.Min.Y; y < b.Max.Y; y++ {
        for x := b.Min.X; x < b.Max.X; x++ {
            px[i].x = x
            px[i].y = y
            i++
        }
    }
    return qz
}

func (qz *quantizer) cluster() {
    
    
    
    
    
    pq := new(queue)
    
    c := &qz.cs[0]
    for i := 1; ; {
        qz.setColorRange(c)
        
        
        if c.chRange > 0 {
            heap.Push(pq, c) 
        }
        
        
        if len(*pq) == 0 {
            qz.cs = qz.cs[:i]
            break
        }
        s := heap.Pop(pq).(*cluster) 
        c = &qz.cs[i]                
        i++
        m := qz.Median(s)
        qz.Split(s, c, m) 
        
        if i == len(qz.cs) {
            break
        }
        qz.setColorRange(s)
        if s.chRange > 0 {
            heap.Push(pq, s) 
        }
    }
}
    
func (q *quantizer) setColorRange(c *cluster) {
    
    var maxR, maxG, maxB uint32
    minR := uint32(math.MaxUint32)
    minG := uint32(math.MaxUint32)
    minB := uint32(math.MaxUint32) 
    for _, p := range c.px {
        r, g, b, _ := q.img.At(p.x, p.y).RGBA()
        if r < minR { 
            minR = r
        }
        if r > maxR {
            maxR = r
        }
        if g < minG {
            minG = g 
        }
        if g > maxG {
            maxG = g
        }
        if b < minB {
            minB = b
        }
        if b > maxB {
            maxB = b
        }
    }
    
    s := gx
    min := minG
    max := maxG
    if maxR-minR > max-min {
        s = rx
        min = minR
        max = maxR
    }
    if maxB-minB > max-min {
        s = bx
        min = minB
        max = maxB
    }
    c.widestCh = s
    c.chRange = max - min 
}

func (q *quantizer) Median(c *cluster) uint32 {
    px := c.px
    ch := q.ch[:len(px)]
    
    switch c.widestCh {
    case rx:
        for i, p := range c.px {
            ch[i], _, _, _ = q.img.At(p.x, p.y).RGBA()
        }
    case gx:
        for i, p := range c.px {
            _, ch[i], _, _ = q.img.At(p.x, p.y).RGBA()
        }
    case bx:
        for i, p := range c.px {
            _, _, ch[i], _ = q.img.At(p.x, p.y).RGBA()
        }
    }
    
    sort.Sort(ch)
    half := len(ch) / 2
    m := ch[half]
    if len(ch)%2 == 0 {
        m = (m + ch[half-1]) / 2
    }
    return m
}

func (q *quantizer) Split(s, c *cluster, m uint32) {
    px := s.px
    var v uint32
    i := 0
    lt := 0
    gt := len(px) - 1
    eq := q.eq[:0] 
    for i <= gt {
        
        r, g, b, _ := q.img.At(px[i].x, px[i].y).RGBA()
        switch s.widestCh {
        case rx:
            v = r
        case gx:
            v = g
        case bx:
            v = b
        } 
        
        switch {
        case v < m:
            px[lt] = px[i]
            lt++
            i++
        case v > m:
            px[gt], px[i] = px[i], px[gt]
            gt--
        default:
            eq = append(eq, px[i])
            i++
        }
    }
    
    if len(eq) > 0 {
        copy(px[lt:], eq) 
        
        
        
        if len(px)-i < lt {
            i = lt
        }
        q.eq = eq 
    }
    
    s.px = px[:i]
    c.px = px[i:]
}   
    
func (qz *quantizer) Paletted() *image.Paletted {
    cp := make(color.Palette, len(qz.cs))
    pi := image.NewPaletted(qz.img.Bounds(), cp)
    for i := range qz.cs {
        px := qz.cs[i].px
        
        var rsum, gsum, bsum int64
        for _, p := range px {
            r, g, b, _ := qz.img.At(p.x, p.y).RGBA()
            rsum += int64(r)
            gsum += int64(g)
            bsum += int64(b)
        } 
        n64 := int64(len(px))
        cp[i] = color.NRGBA64{
            uint16(rsum / n64),
            uint16(gsum / n64),
            uint16(bsum / n64),
            0xffff,
        }
        
        for _, p := range px {
            pi.SetColorIndex(p.x, p.y, uint8(i))
        }
    }
    return pi
}


func (c chValues) Len() int           { return len(c) }
func (c chValues) Less(i, j int) bool { return c[i] < c[j] }
func (c chValues) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }


func (q queue) Len() int { return len(q) }


func (q queue) Less(i, j int) bool {
    return len(q[j].px) < len(q[i].px)
}

func (q queue) Swap(i, j int) {
    q[i], q[j] = q[j], q[i]
}
func (pq *queue) Push(x interface{}) {
    c := x.(*cluster)
    *pq = append(*pq, c)
}
func (pq *queue) Pop() interface{} {
    q := *pq
    n := len(q) - 1
    c := q[n]
    *pq = q[:n]
    return c
}