Search code examples
gogota

group [][]string (2D slice) by row value


I'm working with a 2D slice of string in go and I want to group them by "A" column value but I couldn't figure it out.

I tried to use gota data frame but it also doesn't have group by like what is available in pandas.

    input := [][]string{
        []string{"b", "3", "2.9", "5.3"},
        []string{"a", "4", "5.1", "9.1"},
        []string{"b", "4", "6.0", "5.3"},
        []string{"c", "3", "6.0", "5.5"},
        []string{"a", "2", "7.1", "9.2"},
    }

I want to have output like this.

[[b 3 2.9 5.3 4 6.0 5.3] [a 4 5.1 9.1 2 7.1 9.2] [c 3 6.0 5.5]]


Solution

  • The following group() function utilizes a map to collect input string slices with the same [0] element, and then converts it back to 2D slice. This will do your work:

    func group(input [][]string) (output [][]string) {
        tmp := map[string][]string{}
        for _, slice := range input {
            if len(slice) <= 1 {
                continue
            }
            tmp[slice[0]] = append(tmp[slice[0]], slice[1:]...)
        }
        for k := range tmp {
            v := append([]string{k}, tmp[k]...)
            output = append(output, v)
        }
        return
    }
    
    func main() {
        input := [][]string{
            []string{"b", "3", "2.9", "5.3"},
            []string{"a", "4", "5.1", "9.1"},
            []string{"b", "4", "6.0", "5.3"},
            []string{"c", "3", "6.0", "5.5"},
            []string{"a", "2", "7.1", "9.2"},
        }
        fmt.Println(group(input)) // [[a 4 5.1 9.1 2 7.1 9.2] [c 3 6.0 5.5] [b 3 2.9 5.3 4 6.0 5.3]]
    }
    

    Analysis of the above code is left as exercise. :)